diff --git a/hetu-state-store/src/main/java/io/hetu/core/statestore/hazelcast/HazelcastStateStore.java b/hetu-state-store/src/main/java/io/hetu/core/statestore/hazelcast/HazelcastStateStore.java index 0fc10bac9..e953b6ef6 100644 --- a/hetu-state-store/src/main/java/io/hetu/core/statestore/hazelcast/HazelcastStateStore.java +++ b/hetu-state-store/src/main/java/io/hetu/core/statestore/hazelcast/HazelcastStateStore.java @@ -124,7 +124,7 @@ public class HazelcastStateStore "State collection type: " + type.name() + " not supported"); } collections.putIfAbsent(collectionName, collection); - return collection; + return collections.get(collectionName); } @Override @@ -144,7 +144,7 @@ public class HazelcastStateStore collection = new EncryptedStateMap(collection, createCipherService(encryptionType)); } collections.putIfAbsent(name, collection); - return collection; + return (StateMap) collections.get(name); } @Override diff --git a/presto-main/src/main/java/io/prestosql/dispatcher/DispatchManager.java b/presto-main/src/main/java/io/prestosql/dispatcher/DispatchManager.java index 6ff5fc3a4..45b08ccd0 100644 --- a/presto-main/src/main/java/io/prestosql/dispatcher/DispatchManager.java +++ b/presto-main/src/main/java/io/prestosql/dispatcher/DispatchManager.java @@ -39,7 +39,6 @@ import io.prestosql.spi.QueryId; import io.prestosql.spi.resourcegroups.SelectionContext; import io.prestosql.spi.resourcegroups.SelectionCriteria; import io.prestosql.spi.service.PropertyService; -import io.prestosql.spi.statestore.StateStore; import io.prestosql.statestore.SharedQueryState; import io.prestosql.statestore.StateCacheStore; import io.prestosql.statestore.StateFetcher; @@ -55,22 +54,18 @@ import javax.annotation.PostConstruct; import javax.annotation.PreDestroy; import javax.inject.Inject; -import java.io.IOException; import java.text.SimpleDateFormat; import java.util.Date; import java.util.List; import java.util.Map; import java.util.Optional; import java.util.concurrent.Executor; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.locks.Lock; import java.util.stream.Collectors; import java.util.stream.Stream; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.util.concurrent.Futures.immediateFuture; -import static io.prestosql.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR; import static io.prestosql.spi.StandardErrorCode.QUERY_TEXT_TOO_LARGE; import static io.prestosql.util.StatementUtils.getQueryType; import static io.prestosql.util.StatementUtils.isTransactionControlStatement; @@ -237,17 +232,25 @@ public class DispatchManager boolean queryAdded = queryCreated(dispatchQuery); if (queryAdded && !dispatchQuery.isDone()) { - if (!PropertyService.getBooleanProperty(HetuConstant.MULTI_COORDINATOR_ENABLED)) { - try { - resourceGroupManager.submit(dispatchQuery, selectionContext, queryExecutor); + try { + resourceGroupManager.submit(dispatchQuery, selectionContext, queryExecutor); + + if (PropertyService.getBooleanProperty(HetuConstant.MULTI_COORDINATOR_ENABLED) && stateUpdater != null) { + stateUpdater.registerQuery(StateStoreConstants.QUERY_STATE_COLLECTION_NAME, dispatchQuery); } - catch (Throwable e) { - // dispatch query has already been registered, so just fail it directly - dispatchQuery.fail(e); + + if (LOG.isDebugEnabled()) { + long now = System.currentTimeMillis(); + LOG.debug("query:%s submission started at %s, ended at %s, total time use: %sms", + dispatchQuery.getQueryId(), + new SimpleDateFormat("HH:mm:ss:SSS").format(dispatchQuery.getCreateTime().toDate()), + new SimpleDateFormat("HH:mm:ss:SSS").format(new Date(now)), + now - dispatchQuery.getCreateTime().getMillis()); } } - else { - submitQuerySync(dispatchQuery, selectionContext); + catch (Throwable e) { + // dispatch query has already been registered, so just fail it directly + dispatchQuery.fail(e); } } } @@ -419,56 +422,4 @@ public class DispatchManager stateFetcher.stop(); } } - - // Submit query synchronously to the distributed resource group - private synchronized void submitQuerySync(DispatchQuery dispatchQuery, SelectionContext selectionContext) - throws InterruptedException, PrestoException - { - StateStore stateStore = stateStoreProvider.getStateStore(); - if (stateStore == null) { - LOG.error("StateStore is not loaded yet"); - throw new PrestoException(GENERIC_INTERNAL_ERROR, "Coordinator is not ready to accept queries"); - } - - Lock lock = stateStore.getLock(StateStoreConstants.SUBMIT_QUERY_LOCK_NAME); - // Make sure query submission is synchronized - boolean locked = lock.tryLock(hetuConfig.getQuerySubmitTimeout().toMillis(), TimeUnit.MILLISECONDS); - long start = 0L; - if (locked) { - try { - start = System.currentTimeMillis(); - LOG.debug("Get submit-query-lock, will submit query:%s, at current time milliseconds: %s, at format HH:mm:ss:SSS:%s", - dispatchQuery.getQueryId(), - start, - new SimpleDateFormat("HH:mm:ss:SSS").format(new Date(start))); - stateFetcher.fetchRunningQueryStates(stateStore); - resourceGroupManager.submit(dispatchQuery, selectionContext, queryExecutor); - // Register dispatch query to StateUpdater - if (PropertyService.getBooleanProperty(HetuConstant.MULTI_COORDINATOR_ENABLED) && stateUpdater != null) { - stateUpdater.registerQuery(StateStoreConstants.QUERY_STATE_COLLECTION_NAME, dispatchQuery); - } - stateUpdater.updateStates(); - } - catch (IOException e) { - throw new PrestoException(GENERIC_INTERNAL_ERROR, "Failed to fetch states from or update states to state store: " + e.getMessage()); - } - catch (Throwable e) { - // dispatch query has already been registered, so just fail it directly - dispatchQuery.fail(e); - } - finally { - lock.unlock(); - long end = System.currentTimeMillis(); - LOG.debug("Release submit-query-lock, query:%s, at current time milliseconds: %s, at format HH:mm:ss:SSS:%s, total time use: %s", - dispatchQuery.getQueryId(), - end, - new SimpleDateFormat("HH:mm:ss:SSS").format(new Date(end)), - end - start); - } - } - else { - // TODO maybe just queue the query if the queue size is not a problem - throw new PrestoException(GENERIC_INTERNAL_ERROR, "Coordinator probably too busy at the moment, please try again in a few minutes"); - } - } } diff --git a/presto-main/src/main/java/io/prestosql/execution/resourcegroups/DistributedResourceGroup.java b/presto-main/src/main/java/io/prestosql/execution/resourcegroups/DistributedResourceGroup.java index 3fe8652ee..7fcf14b51 100644 --- a/presto-main/src/main/java/io/prestosql/execution/resourcegroups/DistributedResourceGroup.java +++ b/presto-main/src/main/java/io/prestosql/execution/resourcegroups/DistributedResourceGroup.java @@ -72,7 +72,10 @@ import static java.util.Objects.requireNonNull; * state store and check if query can run or can queue using calculated resource usage * * @since 2019-11-29 + * + * @deprecated , this class is replaced by DistributedResourceGroupTemp */ +@Deprecated @ThreadSafe public class DistributedResourceGroup extends BaseResourceGroup diff --git a/presto-main/src/main/java/io/prestosql/execution/resourcegroups/DistributedResourceGroupAggrStats.java b/presto-main/src/main/java/io/prestosql/execution/resourcegroups/DistributedResourceGroupAggrStats.java new file mode 100644 index 000000000..133dcfefe --- /dev/null +++ b/presto-main/src/main/java/io/prestosql/execution/resourcegroups/DistributedResourceGroupAggrStats.java @@ -0,0 +1,125 @@ +/* + * Copyright (C) 2018-2020. Huawei Technologies Co., Ltd. All rights reserved. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.prestosql.execution.resourcegroups; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import io.prestosql.spi.resourcegroups.ResourceGroupId; + +import java.util.Objects; + +public class DistributedResourceGroupAggrStats +{ + private final ResourceGroupId resourceGroupId; + private final int runningQueries; + private final int queuedQueries; + private final int descendantRunningQueries; + private final int descendantQueuedQueries; + private final long cpuUsageMillis; + private final long cachedMemoryUsageBytes; + + @JsonCreator + public DistributedResourceGroupAggrStats( + @JsonProperty("resourceGroupId") ResourceGroupId resourceGroupId, + @JsonProperty("runningQueries") int runningQueries, + @JsonProperty("queuedQueries") int queuedQueries, + @JsonProperty("descendantRunningQueries") int descendantRunningQueries, + @JsonProperty("descendantQueuedQueries") int descendantQueuedQueries, + @JsonProperty("cpuUsageMillis") long cpuUsageMillis, + @JsonProperty("cachedMemoryUsageBytes") long cachedMemoryUsageBytes) + { + this.runningQueries = runningQueries; + this.queuedQueries = queuedQueries; + this.descendantRunningQueries = descendantRunningQueries; + this.descendantQueuedQueries = descendantQueuedQueries; + this.cpuUsageMillis = cpuUsageMillis; + this.cachedMemoryUsageBytes = cachedMemoryUsageBytes; + this.resourceGroupId = resourceGroupId; + } + + @JsonProperty + public int getRunningQueries() + { + return runningQueries; + } + + @JsonProperty + public int getQueuedQueries() + { + return queuedQueries; + } + + @JsonProperty + public int getDescendantRunningQueries() + { + return descendantRunningQueries; + } + + @JsonProperty + public int getDescendantQueuedQueries() + { + return descendantQueuedQueries; + } + + @JsonProperty + public long getCpuUsageMillis() + { + return cpuUsageMillis; + } + + @JsonProperty + public long getCachedMemoryUsageBytes() + { + return cachedMemoryUsageBytes; + } + + @Override + public boolean equals(Object o) + { + if (this == o) { + return true; + } + if (!(o instanceof DistributedResourceGroupAggrStats)) { + return false; + } + DistributedResourceGroupAggrStats that = (DistributedResourceGroupAggrStats) o; + return resourceGroupId.equals(that.resourceGroupId) + && runningQueries == that.runningQueries + && queuedQueries == that.queuedQueries + && descendantRunningQueries == that.descendantRunningQueries + && descendantQueuedQueries == that.descendantQueuedQueries + && cpuUsageMillis == that.cpuUsageMillis + && cachedMemoryUsageBytes == that.cachedMemoryUsageBytes; + } + + @Override + public int hashCode() + { + return Objects.hash(resourceGroupId, runningQueries, queuedQueries, descendantRunningQueries, descendantQueuedQueries, cpuUsageMillis, cachedMemoryUsageBytes); + } + + @Override + public String toString() + { + return "DistributedResourceGroupAggrStats{" + + "runningQueries=" + runningQueries + + ", queuedQueries=" + queuedQueries + + ", descendantRunningQueries=" + descendantRunningQueries + + ", descendantQueuedQueries=" + descendantQueuedQueries + + ", cpuUsageMillis=" + cpuUsageMillis + + ", cachedMemoryUsageBytes=" + cachedMemoryUsageBytes + + '}'; + } +} diff --git a/presto-main/src/main/java/io/prestosql/execution/resourcegroups/DistributedResourceGroupTemp.java b/presto-main/src/main/java/io/prestosql/execution/resourcegroups/DistributedResourceGroupTemp.java new file mode 100644 index 000000000..9a3b79cbf --- /dev/null +++ b/presto-main/src/main/java/io/prestosql/execution/resourcegroups/DistributedResourceGroupTemp.java @@ -0,0 +1,884 @@ +/* + * Copyright (C) 2018-2020. Huawei Technologies Co., Ltd. All rights reserved. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.prestosql.execution.resourcegroups; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.collect.ImmutableList; +import io.airlift.json.ObjectMapperProvider; +import io.airlift.log.Logger; +import io.airlift.stats.CounterStat; +import io.airlift.units.DataSize; +import io.airlift.units.Duration; +import io.prestosql.execution.ManagedQueryExecution; +import io.prestosql.execution.QueryState; +import io.prestosql.metadata.InternalNode; +import io.prestosql.metadata.InternalNodeManager; +import io.prestosql.server.QueryStateInfo; +import io.prestosql.spi.PrestoException; +import io.prestosql.spi.resourcegroups.KillPolicy; +import io.prestosql.spi.resourcegroups.ResourceGroupId; +import io.prestosql.spi.resourcegroups.SchedulingPolicy; +import io.prestosql.spi.statestore.StateCollection; +import io.prestosql.spi.statestore.StateMap; +import io.prestosql.spi.statestore.StateStore; +import io.prestosql.statestore.SharedQueryState; +import io.prestosql.statestore.SharedResourceGroupState; +import io.prestosql.statestore.StateCacheStore; +import io.prestosql.statestore.StateStoreConstants; +import io.prestosql.utils.DistributedResourceGroupUtils; +import org.joda.time.DateTime; +import org.weakref.jmx.Managed; + +import javax.annotation.concurrent.GuardedBy; +import javax.annotation.concurrent.ThreadSafe; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.Comparator; +import java.util.HashSet; +import java.util.LinkedList; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.PriorityQueue; +import java.util.Queue; +import java.util.Set; +import java.util.concurrent.Executor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.locks.Lock; +import java.util.function.BiConsumer; +import java.util.stream.Collectors; + +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Preconditions.checkState; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.math.LongMath.saturatedAdd; +import static com.google.common.math.LongMath.saturatedMultiply; +import static com.google.common.math.LongMath.saturatedSubtract; +import static io.prestosql.server.QueryStateInfo.createQueryStateInfo; +import static io.prestosql.spi.ErrorType.USER_ERROR; +import static io.prestosql.spi.StandardErrorCode.GENERIC_INSUFFICIENT_RESOURCES; +import static java.util.Objects.requireNonNull; + +/** + * Resource groups form a tree, and all access to a group is guarded by the root of the tree. + * Queries are submitted to leaf groups. Never to intermediate groups. Intermediate groups + * aggregate resource consumption from their children, and may have their own limitations that + * are enforced. + *

+ * Distributed resource group calculates group resource usage using query states from external + * state store and check if query can run or can queue using calculated resource usage + * + * @since 2019-11-29 + *

+ * TODO: After this class is tested(with HA aggregated stats), it will be renamed to 'DistributedResourceGroup' + */ + +@ThreadSafe +public class DistributedResourceGroupTemp + extends BaseResourceGroup +{ + private static final long MILLISECONDS_PER_SECOND = 1000L; + + private static final Logger LOG = Logger.get(DistributedResourceGroupTemp.class); + + // Live data structures + // ==================== + @GuardedBy("root") + private long lastStartMillis; + @GuardedBy("root") + private final CounterStat timeBetweenStartsSec = new CounterStat(); + // Internal time to track last time refreshing stats from state store + private DateTime lastUpdateTime = new DateTime(); + // Last query execution time in current or all the children groups + private Optional lastExecutionTime = Optional.empty(); + // State Store + private StateStore stateStore; + + private InternalNodeManager internalNodeManager; + private static final ObjectMapper MAPPER = new ObjectMapperProvider().get(); + private static final String DASH = "-"; + private static final String RESOURCE_AGGR_STATS = "resourceaggrstats"; + + // Local variables represent value in the current coordinator + @GuardedBy("root") + private Queue localQueuedQueries = new LinkedList<>(); + @GuardedBy("root") + private final Set localRunningQueries = new HashSet<>(); + @GuardedBy("root") + private int localDescendantRunningQueries; + @GuardedBy("root") + private int localDescendantQueuedQueries; + // Memory usage is cached because it changes very rapidly while queries are running, and would be expensive to track continuously + @GuardedBy("root") + private long localCachedMemoryUsageBytes; + @GuardedBy("root") + private long localCpuUsageMillis; + + // Global variables represent aggregated values among all the coordinators + @GuardedBy("root") + private int globalTotalQueuedQueries; + @GuardedBy("root") + private int globalTotalRunningQueries; + @GuardedBy("root") + private int globalDescendantRunningQueries; + @GuardedBy("root") + private int globalDescendantQueuedQueries; + // Memory usage is cached because it changes very rapidly while queries are running, and would be expensive to track continuously + @GuardedBy("root") + private long globalCachedMemoryUsageBytes; + @GuardedBy("root") + private long globalCpuUsageMillis; + + protected DistributedResourceGroupTemp(Optional parent, + String name, + BiConsumer jmxExportListener, + Executor executor, + StateStore stateStore, + InternalNodeManager internalNodeManager) + { + super(parent, name, jmxExportListener, executor); + this.stateStore = requireNonNull(stateStore, "state store is null"); + this.internalNodeManager = requireNonNull(internalNodeManager, "internalNodeManager is null"); + } + + @Override + protected List getAggregatedRunningQueriesInfo() + { + synchronized (root) { + if (subGroups.isEmpty()) { + Optional resourceGroupState = getSharedResourceGroupState(); + if (!resourceGroupState.isPresent()) { + return ImmutableList.of(); + } + else { + return resourceGroupState.get().getRunningQueries().stream() + .map(SharedQueryState::getBasicQueryInfo) + .map(queryInfo -> createQueryStateInfo(queryInfo, Optional.of(id))) + .collect(toImmutableList()); + } + } + + return subGroups.values().stream() + .map(BaseResourceGroup::getAggregatedRunningQueriesInfo) + .flatMap(List::stream) + .collect(toImmutableList()); + } + } + + @Managed + @Override + public int getRunningQueries() + { + synchronized (root) { + return getTotalGlobalRunningQueries() + getGlobalDescendantRunningQueries(); + } + } + + @Managed + @Override + public int getQueuedQueries() + { + synchronized (root) { + refreshGlobalValues(); + return globalTotalQueuedQueries + globalDescendantQueuedQueries; + } + } + + @Override + public void setSoftMemoryLimit(DataSize limit) + { + synchronized (root) { + this.softMemoryLimitBytes = limit.toBytes(); + } + } + + @Override + public void setSoftCpuLimit(Duration limit) + { + synchronized (root) { + if (limit.toMillis() > hardCpuLimitMillis) { + setHardCpuLimit(limit); + } + this.softCpuLimitMillis = limit.toMillis(); + } + } + + @Override + public void setHardCpuLimit(Duration limit) + { + synchronized (root) { + if (limit.toMillis() < softCpuLimitMillis) { + setSoftCpuLimit(limit); + } + this.hardCpuLimitMillis = limit.toMillis(); + } + } + + @Override + public void setSoftConcurrencyLimit(int softConcurrencyLimit) + { + checkArgument(softConcurrencyLimit >= 0, "softConcurrencyLimit is negative"); + synchronized (root) { + this.softConcurrencyLimit = softConcurrencyLimit; + } + } + + @Override + public void setHardReservedConcurrency(int hardReservedConcurrency) + { + checkArgument(hardReservedConcurrency >= 0, "hardReservedConcurrency is negative"); + synchronized (root) { + this.hardReservedConcurrency = hardReservedConcurrency; + } + } + + @Override + public void setSoftReservedMemory(DataSize softReservedMemory) + { + synchronized (root) { + this.softReservedMemory = softReservedMemory.toBytes(); + } + } + + @Managed + @Override + public void setHardConcurrencyLimit(int hardConcurrencyLimit) + { + checkArgument(hardConcurrencyLimit >= 0, "hardConcurrencyLimit is negative"); + synchronized (root) { + this.hardConcurrencyLimit = hardConcurrencyLimit; + } + } + + @Override + public void setSchedulingWeight(int weight) + { + checkArgument(weight > 0, "weight must be positive"); + synchronized (root) { + this.schedulingWeight = weight; + } + } + + @Override + public void setSchedulingPolicy(SchedulingPolicy policy) + { + synchronized (root) { + if (policy == schedulingPolicy) { + return; + } + + // Only Fair scheduling policy is supported for distributed resource group + switch (policy) { + case FAIR: + break; + case WEIGHTED: + case WEIGHTED_FAIR: + case QUERY_PRIORITY: + default: + throw new UnsupportedOperationException("Unsupported scheduling policy: " + policy); + } + schedulingPolicy = policy; + } + } + + @Override + public void setKillPolicy(KillPolicy killPolicy) + { + synchronized (root) { + this.killPolicy = killPolicy; + } + } + + @Override + public DistributedResourceGroupTemp getOrCreateSubGroup(String name) + { + requireNonNull(name, "name is null"); + synchronized (root) { + checkArgument(localRunningQueries.isEmpty() && localQueuedQueries.isEmpty(), "Cannot add sub group to %s while queries are running", id); + if (subGroups.containsKey(name)) { + return (DistributedResourceGroupTemp) subGroups.get(name); + } + DistributedResourceGroupTemp subGroup = new DistributedResourceGroupTemp(Optional.of(this), name, jmxExportListener, executor, stateStore, internalNodeManager); + subGroup.setMemoryMarginPercent(memoryMarginPercent); + subGroup.setQueryProgressMarginPercent(queryProgressMarginPercent); + subGroups.put(name, subGroup); + return subGroup; + } + } + + public Optional getLastExecutionTime() + { + return lastExecutionTime; + } + + /** + * Check availability and try to run the query in current resource group + * + * @param query Query execution contains query states + */ + @Override + public void run(ManagedQueryExecution query) + { + synchronized (root) { + root.internalRefreshStats(); + super.run(query); + } + } + + @Override + protected void enqueueQuery(ManagedQueryExecution query) + { + checkState(Thread.holdsLock(root), "Must hold lock to enqueue a query"); + synchronized (root) { + localQueuedQueries.add(query); + updateLocalValuesToStateStore(); + } + } + + @Override + protected void startInBackground(ManagedQueryExecution query) + { + checkState(Thread.holdsLock(root), "Must hold lock to start a query"); + synchronized (root) { + executor.execute(query::startWaitingForResources); + while (query.getBasicQueryInfo().getState() == QueryState.QUEUED) { + // wait for query to be started + } + localQueuedQueries.remove(query); + localRunningQueries.add(query); + updateLocalValuesToStateStore(); + } + } + + @Override + protected void queryFinished(ManagedQueryExecution query) + { + synchronized (root) { + if (!localRunningQueries.contains(query) && !localQueuedQueries.contains(query)) { + // Query has already been cleaned up + return; + } + + // Only count the CPU time if the query succeeded, or the failure was the fault of the user + if (!query.getErrorCode().isPresent() || query.getErrorCode().get().getType() == USER_ERROR) { + DistributedResourceGroupTemp group = this; + while (group != null) { + group.localCpuUsageMillis = saturatedAdd(group.localCpuUsageMillis, query.getTotalCpuTime().toMillis()); + group = (DistributedResourceGroupTemp) group.parent.orElse(null); + } + } + + if (localRunningQueries.contains(query)) { + localRunningQueries.remove(query); + } + else { + localQueuedQueries.remove(query); + } + updateLocalValuesToStateStore(); + } + } + + /** + * Update internal stats for all the children groups + * Stats are calculated based on cached resource group states from state store + * Stats include: running queries, queued queries, memory usage, cpu usage + */ + @Override + protected void internalRefreshStats() + { + checkState(Thread.holdsLock(root), "Must hold lock to refresh stats"); + synchronized (root) { + if (subGroups.isEmpty()) { + localDescendantRunningQueries = 0; + localDescendantQueuedQueries = 0; + localCachedMemoryUsageBytes = 0; + + for (ManagedQueryExecution query : localRunningQueries) { + localCachedMemoryUsageBytes += query.getTotalMemoryReservation().toBytes(); + } + } + else { + int tempLocalDescendantRunningQueries = 0; + int tempLocalDescendantQueuedQueries = 0; + long tempLocalCachedMemoryUsageBytes = 0L; + long tempLocalCpuUsageMillis = 0L; + + for (BaseResourceGroup group : subGroups()) { + group.internalRefreshStats(); + tempLocalCpuUsageMillis += ((DistributedResourceGroupTemp) group).localCpuUsageMillis; + tempLocalDescendantRunningQueries += ((DistributedResourceGroupTemp) group).localDescendantRunningQueries; + tempLocalDescendantQueuedQueries += ((DistributedResourceGroupTemp) group).localDescendantQueuedQueries; + tempLocalCachedMemoryUsageBytes += ((DistributedResourceGroupTemp) group).localCachedMemoryUsageBytes; + } + + localDescendantRunningQueries = tempLocalDescendantRunningQueries; + localDescendantQueuedQueries = tempLocalDescendantQueuedQueries; + localCachedMemoryUsageBytes = tempLocalCachedMemoryUsageBytes; + localCpuUsageMillis = tempLocalCpuUsageMillis; + } + lastUpdateTime = new DateTime(); + updateLocalValuesToStateStore(); + } + } + + public void internalCancelQuery() + { + checkState(Thread.holdsLock(root), "Must hold lock to check cancel query"); + synchronized (root) { + if (!subGroups.isEmpty()) { + for (BaseResourceGroup group : subGroups()) { + ((DistributedResourceGroupTemp) group).internalCancelQuery(); + } + + return; + } + + long globalCachedMemoryUsageBytes = getGlobalCachedMemoryUsageBytes(); + if (globalCachedMemoryUsageBytes <= softMemoryLimitBytes) { + return; + } + + Optional resourceGroupState = getSharedResourceGroupState(); + if (!resourceGroupState.isPresent()) { + return; + } + + Set globalRunningQueries = resourceGroupState.get().getRunningQueries(); + List sortedQueryList; + + Lock lock = stateStore.getLock(id.toString()); + boolean locked = false; + try { + // If lock is not free, then we return immediately, so no need to refresh after taking lock. + // Before next call of this function, refresh will already happen. + locked = lock.tryLock(); + if (locked) { + switch (killPolicy) { + case HIGH_MEMORY_QUERIES: + double absMemoryMargin = 1 - (double) memoryMarginPercent / 100; + double absQueryProgressMargin = 1 - (double) queryProgressMarginPercent / 100; + + sortedQueryList = globalRunningQueries.stream().sorted((o1, o2) -> { + if (o1.getTotalMemoryReservation().toBytes() < o2.getTotalMemoryReservation().toBytes() * absMemoryMargin + || o2.getTotalMemoryReservation().toBytes() < o1.getTotalMemoryReservation().toBytes() * absMemoryMargin) { + return ((Long) o2.getTotalMemoryReservation().toBytes()).compareTo(o1.getTotalMemoryReservation().toBytes()); + } + + // if memory usage within 10%, then sort based on % of query completion. + // if query progress difference is within 5%, then order will be decided based on memory itself. + if (o1.getQueryProgress().orElse(0) < o2.getQueryProgress().orElse(0) * absQueryProgressMargin + || o2.getQueryProgress().orElse(0) < o1.getQueryProgress().orElse(0) * absQueryProgressMargin) { + return ((Double) o1.getQueryProgress().orElse(0)).compareTo((o2.getQueryProgress().orElse(0))); + } + + return ((Long) o2.getTotalMemoryReservation().toBytes()).compareTo(o1.getTotalMemoryReservation().toBytes()); + }).collect(Collectors.toList()); + break; + case OLDEST_QUERIES: + sortedQueryList = globalRunningQueries.stream().sorted(Comparator.comparing(o -> (o.getExecutionStartTime().get()))).collect(Collectors.toList()); + break; + case RECENT_QUERIES: + sortedQueryList = globalRunningQueries.stream().sorted(Comparator.comparing(o -> (o.getExecutionStartTime().get()), Comparator.reverseOrder())).collect(Collectors.toList()); + break; + case FINISH_PERCENTAGE_QUERIES: + sortedQueryList = globalRunningQueries.stream().sorted(Comparator.comparing(o -> (o.getQueryProgress().orElse(0)))).collect(Collectors.toList()); + break; + case NO_KILL: + //fall through + default: + sortedQueryList = new ArrayList<>(); + } + + long tempGlobalCachedMemoryUsage = globalCachedMemoryUsageBytes; + long tempLocalCachedMemoryUsage = localCachedMemoryUsageBytes; + + // As per the kill policy, top queries are selected across all coordinators but we kill only local queries + // till memory reaches with-in required limit. + // E.g. Suppose Kill policy is HIGH_MEMORY_QUERIES and queries are ordered as below + // Q1, Q10, Q7, Q20, Q25 + // where only Q1 and Q7 are local queries and only combined memory of Q1 and Q10 brings memory within the desired limit. + // So in this case only Q1 will be killed from local coordinator. + for (SharedQueryState query : sortedQueryList) { + for (ManagedQueryExecution localQuery : localRunningQueries) { + if (query.getBasicQueryInfo().getQueryId().equals(localQuery.getBasicQueryInfo().getQueryId())) { + LOG.info("Query " + localQuery.getBasicQueryInfo().getQueryId() + " is getting killed for resource group " + this + " query will be killed with policy " + killPolicy); + localQuery.fail(new PrestoException(GENERIC_INSUFFICIENT_RESOURCES, "Memory consumption " + tempLocalCachedMemoryUsage + " exceed the limit " + softMemoryLimitBytes + "for resource group " + this)); + queryFinished(localQuery); + tempLocalCachedMemoryUsage -= query.getTotalMemoryReservation().toBytes(); + break; + } + } + + tempGlobalCachedMemoryUsage -= query.getTotalMemoryReservation().toBytes(); + if (tempGlobalCachedMemoryUsage <= softMemoryLimitBytes) { + break; + } + } + } + } + catch (RuntimeException e) { + return; + } + finally { + if (locked) { + lock.unlock(); + } + } + } + } + + /** + * Check if there is any queued query in any eligible sub-group that can be started + * + * @return if any query has been started in sub-groups + */ + protected boolean internalStartNext() + { + checkState(Thread.holdsLock(root), "Must hold lock to find next query"); + synchronized (root) { + if (!canRunMore()) { + return false; + } + + // Only start the query if it exists locally + Optional resourceGroupState = getSharedResourceGroupState(); + PriorityQueue globalQueuedQueries = resourceGroupState.isPresent() ? resourceGroupState.get().getQueuedQueries() : new PriorityQueue<>(); + if (!globalQueuedQueries.isEmpty() && !localQueuedQueries.isEmpty()) { + // Get queued query with longest queued time from state cache store. + // Remove it if local queued queries contains it. + SharedQueryState nextQuery = globalQueuedQueries.peek(); + for (ManagedQueryExecution localQuery : localQueuedQueries) { + if (nextQuery.getBasicQueryInfo().getQueryId().equals(localQuery.getBasicQueryInfo().getQueryId())) { + Lock lock = stateStore.getLock(id.toString()); + boolean locked = false; + try { + locked = lock.tryLock(MILLISECONDS_PER_SECOND, TimeUnit.MILLISECONDS); + if (locked) { + // Get the most recent cached state store status and check canRunMore again + // Avoid the race condition that state store is updated by other process + // Make sure queued query start is synchronized + DistributedResourceGroupUtils.mapCachedStates(); + if (canRunMore()) { + startInBackground(localQuery); + return true; + } + } + return false; + } + catch (InterruptedException | RuntimeException e) { + return false; + } + finally { + if (locked) { + lock.unlock(); + } + } + } + } + } + + // Try to find least recently used eligible group + DistributedResourceGroupTemp chosenGroup = findLeastRecentlyExecutedSubgroup(); + if (chosenGroup == null) { + return false; + } + + // Try to start queued query in the group + boolean started = chosenGroup.internalStartNext(); + + long currentTime = System.currentTimeMillis(); + if (lastStartMillis != 0) { + timeBetweenStartsSec.update(Math.max(0, (currentTime - lastStartMillis) / MILLISECONDS_PER_SECOND)); + } + lastStartMillis = currentTime; + + return started; + } + } + + @Override + protected boolean canQueueMore() + { + checkState(Thread.holdsLock(root), "Must hold lock"); + synchronized (root) { + return getQueuedQueries() < maxQueuedQueries; + } + } + + @Override + protected boolean canRunMore() + { + checkState(Thread.holdsLock(root), "Must hold lock"); + synchronized (root) { + refreshGlobalValues(); + if (globalCpuUsageMillis >= hardCpuLimitMillis) { + return false; + } + int hardConcurrencyLimit = this.hardConcurrencyLimit; + return hasCapacity(globalTotalRunningQueries, adjustHardConcurrency(hardConcurrencyLimit, globalCpuUsageMillis)); + } + } + + private boolean hasCapacity(int numRunningQueries, int hardConcurrencyLimit) + { + if (numRunningQueries + globalDescendantRunningQueries >= hardConcurrencyLimit || + globalCachedMemoryUsageBytes > softMemoryLimitBytes) { + return false; + } + if (parent.isPresent()) { + // Check if hardConcurrencyLimit of the parent is reached with reserved concurrency for each peer group + if (numRunningQueries + globalDescendantRunningQueries >= hardReservedConcurrency) { + int peerTotalQuerySize = 0; + for (DistributedResourceGroupTemp group : (Collection) parent.get().subGroups()) { + group.refreshGlobalValues(); + peerTotalQuerySize += Math.max(group.globalTotalRunningQueries + group.globalDescendantRunningQueries, group.hardReservedConcurrency); + } + if (parent.get().hardConcurrencyLimit <= peerTotalQuerySize) { + return false; + } + } + // Check if the softMemoryLimit of parent is reached with reserved memory for each peer group + if (globalCachedMemoryUsageBytes >= softReservedMemory) { + long peerGroupTotalUsage = 0L; + for (DistributedResourceGroupTemp group : (Collection) parent.get().subGroups()) { + peerGroupTotalUsage += Math.max(group.globalCachedMemoryUsageBytes, group.softReservedMemory); + } + if (parent.get().softMemoryLimitBytes <= peerGroupTotalUsage) { + LOG.debug("No capacity to run more queries in the resource group: %s, with following reasons: \n" + + "cachedMemoryUsageBytes:%s >= softReservedMemory:%s and \n" + + "softMemoryLimitBytes:%s <= peerGroupTotalUsage:%s", + parent.get().id, globalCachedMemoryUsageBytes, softReservedMemory, softMemoryLimitBytes, peerGroupTotalUsage); + return false; + } + } + } + return true; + } + + protected void internalGenerateCpuQuota(long elapsedSeconds) + { + checkState(Thread.holdsLock(root), "Must hold lock to generate cpu quota"); + synchronized (root) { + // Bug fix in cases when the CPU usage is initially higher than the cpu quota, but after a while, + // when the CPU usage is then lower than the CPU quota, queued queries are still not run. This fix was + long newQuota = saturatedMultiply(elapsedSeconds, cpuQuotaGenerationMillisPerSecond); + localCpuUsageMillis = saturatedSubtract(localCpuUsageMillis, newQuota); + + if (localCpuUsageMillis < 0 || localCpuUsageMillis == Long.MAX_VALUE) { + localCpuUsageMillis = 0; + } + updateLocalValuesToStateStore(); + for (BaseResourceGroup group : subGroups.values()) { + ((InternalResourceGroup) group).internalGenerateCpuQuota(elapsedSeconds); + } + } + } + + /** + * Return subgroups of current resource group + * + * @return collection contains all the subgroups + */ + @Override + public Collection subGroups() + { + synchronized (root) { + return subGroups.values(); + } + } + + /** + * Get cached global resource group state from state store for current resource group + * + * @return SharedResourceGroupState if exists in state store + */ + private Optional getSharedResourceGroupState() + { + Map cachedStates = StateCacheStore.get().getCachedStates(StateStoreConstants.RESOURCE_GROUP_STATE_COLLECTION_NAME); + if (cachedStates == null || !cachedStates.containsKey(id)) { + return Optional.empty(); + } + return Optional.ofNullable(cachedStates.get(id)); + } + + /** + * Find the sub group that's least recently used that has queued queries and can run more + * + * @return The chosen group or null if no eligible group + */ + private DistributedResourceGroupTemp findLeastRecentlyExecutedSubgroup() + { + List eligibleGroups = subGroups().stream() + .filter(group -> group.getQueuedQueries() > 0 && group.canRunMore()) + .sorted(Comparator.comparing(group -> group.getId().toString())) + .map(DistributedResourceGroupTemp.class::cast) + .collect(Collectors.toList()); + + DateTime leastRecentlyExecutionTime = null; + DistributedResourceGroupTemp chosenGroup = null; + for (DistributedResourceGroupTemp group : eligibleGroups) { + // If the group has never executed any query just use it + if (!group.getLastExecutionTime().isPresent()) { + return group; + } + + if (leastRecentlyExecutionTime == null || group.getLastExecutionTime().get().isBefore(leastRecentlyExecutionTime)) { + leastRecentlyExecutionTime = group.getLastExecutionTime().get(); + chosenGroup = group; + } + } + + return chosenGroup; + } + + /** + * Periodically check and start queries in queue + */ + @Override + public synchronized void processQueuedQueries() + { + internalRefreshStats(); + internalCancelQuery(); + while (internalStartNext()) { + // start all the queries we can + } + } + + @Override + public synchronized void generateCpuQuota(long elapsedSeconds) + { + if (elapsedSeconds > 0) { + internalGenerateCpuQuota(elapsedSeconds); + } + } + + @Override + public long getCachedMemoryUsageBytes() + { + return getGlobalCachedMemoryUsageBytes(); + } + + public int getGlobalDescendantRunningQueries() + { + synchronized (root) { + refreshGlobalValues(); + return globalDescendantRunningQueries; + } + } + + public int getGlobalDescendantQueuedQueries() + { + synchronized (root) { + refreshGlobalValues(); + return globalDescendantQueuedQueries; + } + } + + public long getGlobalCachedMemoryUsageBytes() + { + synchronized (root) { + refreshGlobalValues(); + return globalCachedMemoryUsageBytes; + } + } + + public long getGlobalCpuUsageMillis() + { + synchronized (root) { + refreshGlobalValues(); + return globalCpuUsageMillis; + } + } + + public int getTotalGlobalQueuedQueries() + { + synchronized (root) { + refreshGlobalValues(); + return globalTotalQueuedQueries; + } + } + + public int getTotalGlobalRunningQueries() + { + synchronized (root) { + refreshGlobalValues(); + return globalTotalRunningQueries; + } + } + + private String createCoordinatorCollectionName(InternalNode coordinator) + { + return coordinator.getHostAndPort() + DASH + RESOURCE_AGGR_STATS; + } + + private void updateLocalValuesToStateStore() + { + synchronized (root) { + try { + StateMap resourceGroupMap = ((StateMap) stateStore.getOrCreateStateCollection(createCoordinatorCollectionName(internalNodeManager.getCurrentNode()), StateCollection.Type.MAP)); + DistributedResourceGroupAggrStats groupAggrStats = new DistributedResourceGroupAggrStats( + getId(), + localRunningQueries.size(), + localQueuedQueries.size(), + localDescendantRunningQueries, + localDescendantQueuedQueries, + localCpuUsageMillis, + localCachedMemoryUsageBytes); + String json = MAPPER.writeValueAsString(groupAggrStats); + resourceGroupMap.put(getId().toString(), json); + } + catch (JsonProcessingException e) { + throw new RuntimeException(String.format("Error updating resource group state with group id = %s, caused by ObjectMapper: %s", id, e.getMessage())); + } + } + } + + private void refreshGlobalValues() + { + synchronized (root) { + globalTotalRunningQueries = localRunningQueries.size(); + globalTotalQueuedQueries = localQueuedQueries.size(); + globalDescendantQueuedQueries = localDescendantQueuedQueries; + globalDescendantRunningQueries = localDescendantRunningQueries; + globalCachedMemoryUsageBytes = localCachedMemoryUsageBytes; + globalCpuUsageMillis = localCpuUsageMillis; + + internalNodeManager.refreshNodes(); + try { + for (InternalNode coordinator : internalNodeManager.getCoordinators()) { + if (coordinator.equals(internalNodeManager.getCurrentNode())) { + continue; + } + StateMap resourceGroupMap = ((StateMap) stateStore.getOrCreateStateCollection(createCoordinatorCollectionName(coordinator), StateCollection.Type.MAP)); + DistributedResourceGroupAggrStats groupAggrStats = resourceGroupMap.containsKey(getId().toString()) ? MAPPER.readerFor(DistributedResourceGroupAggrStats.class) + .readValue(resourceGroupMap.get(getId().toString())) : null; + if (groupAggrStats != null) { + globalTotalRunningQueries += groupAggrStats.getRunningQueries(); + globalTotalQueuedQueries += groupAggrStats.getQueuedQueries(); + globalDescendantQueuedQueries += groupAggrStats.getDescendantQueuedQueries(); + globalDescendantRunningQueries += groupAggrStats.getDescendantRunningQueries(); + globalCachedMemoryUsageBytes += groupAggrStats.getCachedMemoryUsageBytes(); + globalCpuUsageMillis += groupAggrStats.getCachedMemoryUsageBytes(); + } + } + } + catch (JsonProcessingException e) { + throw new RuntimeException(String.format("Error fetching resource group state with group id = %s, caused by ObjectMapper: %s", id, e.getMessage())); + } + } + } +} diff --git a/presto-main/src/main/java/io/prestosql/execution/resourcegroups/InternalResourceGroupManager.java b/presto-main/src/main/java/io/prestosql/execution/resourcegroups/InternalResourceGroupManager.java index 6527fb091..4ddfad4fb 100644 --- a/presto-main/src/main/java/io/prestosql/execution/resourcegroups/InternalResourceGroupManager.java +++ b/presto-main/src/main/java/io/prestosql/execution/resourcegroups/InternalResourceGroupManager.java @@ -18,6 +18,7 @@ import com.google.common.collect.ImmutableMap; import io.airlift.log.Logger; import io.airlift.node.NodeInfo; import io.prestosql.execution.ManagedQueryExecution; +import io.prestosql.metadata.InternalNodeManager; import io.prestosql.server.ResourceGroupInfo; import io.prestosql.spi.PrestoException; import io.prestosql.spi.memory.ClusterMemoryPoolManager; @@ -62,6 +63,7 @@ import static com.google.common.base.Preconditions.checkState; import static com.google.common.base.Strings.isNullOrEmpty; import static io.airlift.concurrent.Threads.daemonThreadsNamed; import static io.airlift.configuration.ConfigurationLoader.loadPropertiesFrom; +import static io.prestosql.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR; import static io.prestosql.spi.StandardErrorCode.QUERY_REJECTED; import static java.lang.String.format; import static java.util.Objects.requireNonNull; @@ -79,6 +81,7 @@ public final class InternalResourceGroupManager private static final String RESOURCE_GROUP_QUERY_PROGRESS_MARGIN_PERCENT = "resource-groups.query-progress-margin-percent"; // default status refresh interval private static final long DEFAULT_STATUS_REFRESH_INTERVAL = 1L; + private static final long MILLISECONDS_PER_TEN_SECONDS = 10000L; private final ScheduledExecutorService refreshExecutor = newSingleThreadScheduledExecutor(daemonThreadsNamed("ResourceGroupManager")); private final List rootGroups = new CopyOnWriteArrayList<>(); @@ -96,6 +99,7 @@ public final class InternalResourceGroupManager private final StateStoreProvider stateStoreProvider; private int memoryMarginPercent; private int queryProgressMarginPercent; + private InternalNodeManager internalNodeManager; @Inject public InternalResourceGroupManager(LegacyResourceGroupConfigurationManager legacyManager, @@ -103,7 +107,8 @@ public final class InternalResourceGroupManager ClusterMemoryPoolManager memoryPoolManager, NodeInfo nodeInfo, MBeanExporter exporter, - HetuConfig hetuConfig) + HetuConfig hetuConfig, + InternalNodeManager internalNodeManager) { this.exporter = requireNonNull(exporter, "exporter is null"); this.configurationManagerContext = new ResourceGroupConfigurationManagerContextInstance(memoryPoolManager, nodeInfo.getEnvironment()); @@ -115,6 +120,7 @@ public final class InternalResourceGroupManager this.stateStoreProvider = requireNonNull(stateStoreProvider, "stateStoreProvider is null"); this.memoryMarginPercent = 10; this.queryProgressMarginPercent = 5; + this.internalNodeManager = internalNodeManager; } @Override @@ -139,8 +145,35 @@ public final class InternalResourceGroupManager // Update shared resource group states before submitting new query if (isMultiCoordinatorEnabled) { DistributedResourceGroupUtils.mapCachedStates(); + BaseResourceGroup currentRoot = groups.get(selectionContext.getResourceGroupId().getRoot()); + checkState(currentRoot != null, "currentRoot should not be null"); + synchronized (currentRoot) { + Lock lock = stateStoreProvider.getStateStore().getLock(selectionContext.getResourceGroupId().toString()); + boolean locked = false; + try { + locked = Thread.holdsLock(lock) || lock.tryLock(MILLISECONDS_PER_TEN_SECONDS, TimeUnit.MILLISECONDS); + if (locked) { + groups.get(selectionContext.getResourceGroupId()).run(queryExecution); + } + else { + throw new PrestoException(GENERIC_INTERNAL_ERROR, + String.format("Query: %s submitted failed! Coordinator probably too busy, please try again later", + queryExecution.getBasicQueryInfo().getQueryId())); + } + } + catch (InterruptedException e) { + throw new RuntimeException(e); + } + finally { + if (locked) { + lock.unlock(); + } + } + } + } + else { + groups.get(selectionContext.getResourceGroupId()).run(queryExecution); } - groups.get(selectionContext.getResourceGroupId()).run(queryExecution); } @Override @@ -219,32 +252,7 @@ public final class InternalResourceGroupManager private void refreshAndStartQueries() { - //single coordinator - if (!isMultiCoordinatorEnabled) { - long nanoTime = System.nanoTime(); - long elapsedSeconds = NANOSECONDS.toSeconds(nanoTime - lastCpuQuotaGenerationNanos.get()); - if (elapsedSeconds > 0) { - // Only advance our clock on second boundaries to avoid calling generateCpuQuota() too frequently, and because it would be a no-op for zero seconds. - lastCpuQuotaGenerationNanos.addAndGet(elapsedSeconds * 1_000_000_000L); - } - else if (elapsedSeconds < 0) { - // nano time has overflowed - lastCpuQuotaGenerationNanos.set(nanoTime); - } - for (BaseResourceGroup group : rootGroups) { - try { - if (elapsedSeconds > 0) { - group.generateCpuQuota(elapsedSeconds); - } - } - catch (RuntimeException e) { - log.error(e, "Exception while generation cpu quota for %s", group); - } - } - } - - //multiple coordinator - else { + if (isMultiCoordinatorEnabled) { try { DistributedResourceGroupUtils.mapCachedStates(); } @@ -274,12 +282,25 @@ public final class InternalResourceGroupManager } //for both single and multiple coordinator + long nanoTime = System.nanoTime(); + long elapsedSeconds = NANOSECONDS.toSeconds(nanoTime - lastCpuQuotaGenerationNanos.get()); + if (elapsedSeconds > 0) { + // Only advance our clock on second boundaries to avoid calling generateCpuQuota() too frequently, and because it would be a no-op for zero seconds. + lastCpuQuotaGenerationNanos.addAndGet(elapsedSeconds * 1_000_000_000L); + } + else if (elapsedSeconds < 0) { + // nano time has overflowed + lastCpuQuotaGenerationNanos.set(nanoTime); + } for (BaseResourceGroup group : rootGroups) { try { + if (elapsedSeconds > 0) { + group.generateCpuQuota(elapsedSeconds); + } group.processQueuedQueries(); } catch (RuntimeException e) { - log.error(e, "Exception while processing queued queries for %s", group); + log.error(e, "Exception while refreshing for group %s", group); } } } @@ -382,7 +403,7 @@ public final class InternalResourceGroupManager private BaseResourceGroup createNewRootGroup(String name, Executor executor) { if (isMultiCoordinatorEnabled) { - return new DistributedResourceGroup(Optional.empty(), name, this::exportGroup, executor, stateStoreProvider.getStateStore()); + return new DistributedResourceGroupTemp(Optional.empty(), name, this::exportGroup, executor, stateStoreProvider.getStateStore(), internalNodeManager); } else { return new InternalResourceGroup(Optional.empty(), name, this::exportGroup, executor); diff --git a/presto-main/src/test/java/io/prestosql/execution/resourcegroups/TestDistributedResourceGroup.java b/presto-main/src/test/java/io/prestosql/execution/resourcegroups/TestDistributedResourceGroup.java index cfb19825b..34bfba8f4 100644 --- a/presto-main/src/test/java/io/prestosql/execution/resourcegroups/TestDistributedResourceGroup.java +++ b/presto-main/src/test/java/io/prestosql/execution/resourcegroups/TestDistributedResourceGroup.java @@ -14,15 +14,22 @@ */ package io.prestosql.execution.resourcegroups; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.collect.ImmutableSet; +import io.airlift.json.ObjectMapperProvider; import io.airlift.units.DataSize; import io.airlift.units.Duration; +import io.prestosql.client.NodeVersion; import io.prestosql.execution.MockManagedQueryExecution; +import io.prestosql.metadata.InternalNode; +import io.prestosql.metadata.InternalNodeManager; import io.prestosql.server.QueryStateInfo; import io.prestosql.server.ResourceGroupInfo; import io.prestosql.spi.resourcegroups.KillPolicy; import io.prestosql.spi.resourcegroups.ResourceGroupId; import io.prestosql.spi.statestore.StateStore; +import io.prestosql.statestore.MockStateMap; import io.prestosql.statestore.SharedQueryState; import io.prestosql.statestore.SharedResourceGroupState; import io.prestosql.statestore.StateCacheStore; @@ -34,6 +41,7 @@ import org.mockito.internal.stubbing.answers.Returns; import org.testng.annotations.BeforeClass; import org.testng.annotations.Test; +import java.net.URI; import java.util.HashMap; import java.util.HashSet; import java.util.Iterator; @@ -57,6 +65,7 @@ import static io.prestosql.spi.resourcegroups.SchedulingPolicy.WEIGHTED; import static io.prestosql.spi.resourcegroups.SchedulingPolicy.WEIGHTED_FAIR; import static java.util.concurrent.TimeUnit.MILLISECONDS; import static java.util.concurrent.TimeUnit.SECONDS; +import static org.mockito.Matchers.anyObject; import static org.mockito.Matchers.anyString; import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; @@ -72,17 +81,53 @@ public class TestDistributedResourceGroup // To ensure that test cases are run sequentially private final Object lock = new Object(); private StateStore statestore; + private InternalNodeManager internalNodeManager; private static final DataSize ONE_BYTE = new DataSize(1, BYTE); private static final DataSize FIVE_BYTE = new DataSize(5, BYTE); private static final DataSize TEN_BYTE = new DataSize(10, BYTE); private static final DataSize ONE_MEGABYTE = new DataSize(1, MEGABYTE); private static final DataSize ONE_GIGABYTE = new DataSize(1, GIGABYTE); + private static final ObjectMapper MAPPER = new ObjectMapperProvider().get(); @BeforeClass public void setup() { statestore = Mockito.mock(StateStore.class); when(statestore.getLock(anyString())).then(new Returns(new ReentrantLock())); + internalNodeManager = Mockito.mock(InternalNodeManager.class); + } + + @Test + public void testStateStoreFetchAndUpdate() + throws JsonProcessingException + { + synchronized (lock) { + DistributedResourceGroupTemp root = new DistributedResourceGroupTemp(Optional.empty(), "root", (group, export) -> {}, directExecutor(), statestore, internalNodeManager); + resourceGroupBasicSetUp(root, ONE_MEGABYTE, 1, 1); + MockManagedQueryExecution query1 = new MockManagedQueryExecution(100); + query1.setResourceGroupId(root.getId()); + MockManagedQueryExecution query2 = new MockManagedQueryExecution(0); + query2.setResourceGroupId(root.getId()); + Map mockMap = new HashMap<>(); + MockStateMap mockStateMap = new MockStateMap<>("127.0.0.1-resourceaggrstats", mockMap); + when(statestore.getOrCreateStateCollection(anyString(), anyObject())).thenReturn(mockStateMap); + when(internalNodeManager.getCurrentNode()) + .thenReturn(new InternalNode("node1", URI.create("local://127.0.0.1"), NodeVersion.UNKNOWN, true)); + when(internalNodeManager.getCoordinators()) + .thenReturn(ImmutableSet.of(new InternalNode("node1", URI.create("local://127.0.0.1"), NodeVersion.UNKNOWN, true))); + root.run(query1); + DistributedResourceGroupAggrStats rootStats = MAPPER.readerFor(DistributedResourceGroupAggrStats.class) + .readValue(mockMap.get("root")); + assertEquals(rootStats.getRunningQueries(), 1); + root.run(query2); + rootStats = MAPPER.readerFor(DistributedResourceGroupAggrStats.class) + .readValue(mockMap.get("root")); + assertEquals(rootStats.getQueuedQueries(), 1); + assertEquals(rootStats.getCachedMemoryUsageBytes(), 100); + when(internalNodeManager.getCurrentNode()) + .thenReturn(new InternalNode("node2", URI.create("local://127.0.0.2"), NodeVersion.UNKNOWN, true)); + assertEquals(root.getGlobalCachedMemoryUsageBytes(), 200); + } } @Test @@ -704,7 +749,8 @@ public class TestDistributedResourceGroup } @Test - public void testQueryKillFinishPercent() throws InterruptedException + public void testQueryKillFinishPercent() + throws InterruptedException { synchronized (lock) { DistributedResourceGroup root = new DistributedResourceGroup(Optional.empty(), "root", (group, export) -> {}, directExecutor(), statestore); @@ -754,6 +800,13 @@ public class TestDistributedResourceGroup StateCacheStore.get().setCachedStates(StateStoreConstants.RESOURCE_GROUP_STATE_COLLECTION_NAME, resourceGroupStates); } + private void resourceGroupBasicSetUp(DistributedResourceGroupTemp group, DataSize softMemoryLimit, int hardConcurrencyLimit, int maxQueued) + { + group.setSoftMemoryLimit(softMemoryLimit); + group.setMaxQueuedQueries(maxQueued); + group.setHardConcurrencyLimit(hardConcurrencyLimit); + } + // Set up mandatory fields for DistributedResourceGroup private void resourceGroupBasicSetUp(DistributedResourceGroup group, DataSize softMemoryLimit, int hardConcurrencyLimit, int maxQueued) {