diff --git a/presto-main/src/main/java/com/facebook/presto/dispatcher/CoordinatorLocation.java b/presto-main/src/main/java/com/facebook/presto/dispatcher/CoordinatorLocation.java new file mode 100644 index 0000000000000..83c606210ce46 --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/dispatcher/CoordinatorLocation.java @@ -0,0 +1,23 @@ +/* + * 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 com.facebook.presto.dispatcher; + +import javax.ws.rs.core.UriInfo; + +import java.net.URI; + +public interface CoordinatorLocation +{ + URI getUri(UriInfo uriInfo, String xForwardedProto); +} diff --git a/presto-main/src/main/java/com/facebook/presto/dispatcher/DispatchExecutor.java b/presto-main/src/main/java/com/facebook/presto/dispatcher/DispatchExecutor.java new file mode 100644 index 0000000000000..4eb97a1822ab3 --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/dispatcher/DispatchExecutor.java @@ -0,0 +1,112 @@ +/* + * 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 com.facebook.presto.dispatcher; + +import com.facebook.presto.execution.QueryManagerConfig; +import com.google.common.io.Closer; +import com.google.common.util.concurrent.ListeningExecutorService; +import com.google.common.util.concurrent.ListeningScheduledExecutorService; +import io.airlift.concurrent.ThreadPoolExecutorMBean; +import org.weakref.jmx.Flatten; +import org.weakref.jmx.Managed; +import org.weakref.jmx.Nested; + +import javax.annotation.PreDestroy; +import javax.inject.Inject; + +import java.util.concurrent.ExecutorService; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ThreadPoolExecutor; + +import static com.google.common.util.concurrent.MoreExecutors.listeningDecorator; +import static io.airlift.concurrent.Threads.daemonThreadsNamed; +import static java.util.Objects.requireNonNull; +import static java.util.concurrent.Executors.newCachedThreadPool; +import static java.util.concurrent.Executors.newScheduledThreadPool; + +public class DispatchExecutor +{ + private final Closer closer = Closer.create(); + + private final ListeningExecutorService executor; + private final ListeningScheduledExecutorService scheduledExecutor; + + private final DispatchExecutorMBeans mbeans; + + @Inject + public DispatchExecutor(QueryManagerConfig config) + { + ExecutorService coreExecutor = newCachedThreadPool(daemonThreadsNamed("dispatcher-query-%s")); + closer.register(coreExecutor::shutdownNow); + executor = listeningDecorator(coreExecutor); + + ScheduledExecutorService coreScheduledExecutor = newScheduledThreadPool(config.getQueryManagerExecutorPoolSize(), daemonThreadsNamed("dispatch-executor-%s")); + closer.register(coreScheduledExecutor::shutdownNow); + scheduledExecutor = listeningDecorator(coreScheduledExecutor); + + mbeans = new DispatchExecutorMBeans(coreExecutor, coreScheduledExecutor); + } + + public ListeningExecutorService getExecutor() + { + return executor; + } + + public ListeningScheduledExecutorService getScheduledExecutor() + { + return scheduledExecutor; + } + + @Managed + @Flatten + public DispatchExecutorMBeans getMbeans() + { + return mbeans; + } + + @PreDestroy + public void shutdown() + throws Exception + { + closer.close(); + } + + public class DispatchExecutorMBeans + { + private final ThreadPoolExecutorMBean executor; + private final ThreadPoolExecutorMBean scheduledExecutor; + + public DispatchExecutorMBeans(ExecutorService coreExecutor, ScheduledExecutorService coreScheduledExecutor) + { + requireNonNull(coreExecutor, "coreExecutor is null"); + requireNonNull(coreScheduledExecutor, "coreScheduledExecutor is null"); + executor = new ThreadPoolExecutorMBean((ThreadPoolExecutor) coreExecutor); + scheduledExecutor = new ThreadPoolExecutorMBean((ThreadPoolExecutor) coreScheduledExecutor); + } + + @Managed + @Nested + public ThreadPoolExecutorMBean getExecutor() + { + return executor; + } + + @Managed + @Nested + public ThreadPoolExecutorMBean getScheduledExecutor() + { + return scheduledExecutor; + } + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/dispatcher/DispatchInfo.java b/presto-main/src/main/java/com/facebook/presto/dispatcher/DispatchInfo.java new file mode 100644 index 0000000000000..4700765798ad8 --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/dispatcher/DispatchInfo.java @@ -0,0 +1,74 @@ +/* + * 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 com.facebook.presto.dispatcher; + +import com.facebook.presto.execution.ExecutionFailureInfo; +import io.airlift.units.Duration; + +import java.util.Optional; + +import static java.util.Objects.requireNonNull; + +public class DispatchInfo +{ + private final Optional coordinatorLocation; + private final Optional failureInfo; + private final Duration elapsedTime; + private final Duration queuedTime; + + public static DispatchInfo queued(Duration elapsedTime, Duration queuedTime) + { + return new DispatchInfo(Optional.empty(), Optional.empty(), elapsedTime, queuedTime); + } + + public static DispatchInfo dispatched(CoordinatorLocation coordinatorLocation, Duration elapsedTime, Duration queuedTime) + { + requireNonNull(coordinatorLocation, "coordinatorLocation is null"); + return new DispatchInfo(Optional.of(coordinatorLocation), Optional.empty(), elapsedTime, queuedTime); + } + + public static DispatchInfo failed(ExecutionFailureInfo failureInfo, Duration elapsedTime, Duration queuedTime) + { + requireNonNull(failureInfo, "coordinatorLocation is null"); + return new DispatchInfo(Optional.empty(), Optional.of(failureInfo), elapsedTime, queuedTime); + } + + private DispatchInfo(Optional coordinatorLocation, Optional failureInfo, Duration elapsedTime, Duration queuedTime) + { + this.coordinatorLocation = requireNonNull(coordinatorLocation, "coordinatorLocation is null"); + this.failureInfo = requireNonNull(failureInfo, "failureInfo is null"); + this.elapsedTime = requireNonNull(elapsedTime, "elapsedTime is null"); + this.queuedTime = requireNonNull(queuedTime, "queuedTime is null"); + } + + public Optional getCoordinatorLocation() + { + return coordinatorLocation; + } + + public Optional getFailureInfo() + { + return failureInfo; + } + + public Duration getElapsedTime() + { + return elapsedTime; + } + + public Duration getQueuedTime() + { + return queuedTime; + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/dispatcher/DispatchManager.java b/presto-main/src/main/java/com/facebook/presto/dispatcher/DispatchManager.java new file mode 100644 index 0000000000000..bcc0403937f1b --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/dispatcher/DispatchManager.java @@ -0,0 +1,293 @@ +/* + * 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 com.facebook.presto.dispatcher; + +import com.facebook.presto.Session; +import com.facebook.presto.execution.QueryIdGenerator; +import com.facebook.presto.execution.QueryInfo; +import com.facebook.presto.execution.QueryManagerConfig; +import com.facebook.presto.execution.QueryManagerStats; +import com.facebook.presto.execution.QueryPreparer; +import com.facebook.presto.execution.QueryPreparer.PreparedQuery; +import com.facebook.presto.execution.QueryTracker; +import com.facebook.presto.execution.resourceGroups.ResourceGroupManager; +import com.facebook.presto.execution.warnings.WarningCollector; +import com.facebook.presto.execution.warnings.WarningCollectorFactory; +import com.facebook.presto.metadata.SessionPropertyManager; +import com.facebook.presto.security.AccessControl; +import com.facebook.presto.server.BasicQueryInfo; +import com.facebook.presto.server.SessionContext; +import com.facebook.presto.server.SessionPropertyDefaults; +import com.facebook.presto.server.SessionSupplier; +import com.facebook.presto.spi.PrestoException; +import com.facebook.presto.spi.QueryId; +import com.facebook.presto.spi.resourceGroups.QueryType; +import com.facebook.presto.spi.resourceGroups.SelectionContext; +import com.facebook.presto.spi.resourceGroups.SelectionCriteria; +import com.facebook.presto.transaction.TransactionManager; +import com.google.common.util.concurrent.AbstractFuture; +import com.google.common.util.concurrent.ListenableFuture; +import org.weakref.jmx.Flatten; +import org.weakref.jmx.Managed; + +import javax.inject.Inject; + +import java.util.List; +import java.util.Optional; +import java.util.concurrent.Executor; + +import static com.facebook.presto.spi.StandardErrorCode.QUERY_TEXT_TOO_LARGE; +import static com.facebook.presto.util.StatementUtils.getQueryType; +import static com.facebook.presto.util.StatementUtils.isTransactionControlStatement; +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 java.lang.String.format; +import static java.util.Objects.requireNonNull; + +public class DispatchManager +{ + private final QueryIdGenerator queryIdGenerator; + private final QueryPreparer queryPreparer; + private final ResourceGroupManager resourceGroupManager; + private final WarningCollector warningCollector; + private final DispatchQueryFactory dispatchQueryFactory; + private final FailedDispatchQueryFactory failedDispatchQueryFactory; + private final TransactionManager transactionManager; + private final AccessControl accessControl; + private final SessionSupplier sessionSupplier; + private final SessionPropertyDefaults sessionPropertyDefaults; + + private final int maxQueryLength; + + private final Executor queryExecutor; + + private final QueryTracker queryTracker; + + private final QueryManagerStats stats = new QueryManagerStats(); + + @Inject + public DispatchManager( + QueryIdGenerator queryIdGenerator, + QueryPreparer queryPreparer, + @SuppressWarnings("rawtypes") ResourceGroupManager resourceGroupManager, + WarningCollectorFactory warningCollectorFactory, + DispatchQueryFactory dispatchQueryFactory, + FailedDispatchQueryFactory failedDispatchQueryFactory, + TransactionManager transactionManager, + AccessControl accessControl, + SessionSupplier sessionSupplier, + SessionPropertyDefaults sessionPropertyDefaults, + QueryManagerConfig queryManagerConfig, + DispatchExecutor dispatchExecutor) + { + this.queryIdGenerator = requireNonNull(queryIdGenerator, "queryIdGenerator is null"); + this.queryPreparer = requireNonNull(queryPreparer, "queryPreparer is null"); + this.resourceGroupManager = requireNonNull(resourceGroupManager, "resourceGroupManager is null"); + this.warningCollector = warningCollectorFactory.create(); + this.dispatchQueryFactory = requireNonNull(dispatchQueryFactory, "dispatchQueryFactory is null"); + this.failedDispatchQueryFactory = requireNonNull(failedDispatchQueryFactory, "failedDispatchQueryFactory is null"); + this.transactionManager = requireNonNull(transactionManager, "transactionManager is null"); + this.accessControl = requireNonNull(accessControl, "accessControl is null"); + this.sessionSupplier = requireNonNull(sessionSupplier, "sessionSupplier is null"); + this.sessionPropertyDefaults = requireNonNull(sessionPropertyDefaults, "sessionPropertyDefaults is null"); + + this.maxQueryLength = queryManagerConfig.getMaxQueryLength(); + + this.queryExecutor = requireNonNull(dispatchExecutor, "dispatchExecutor is null").getExecutor(); + + this.queryTracker = new QueryTracker<>(queryManagerConfig, dispatchExecutor.getScheduledExecutor()); + } + + @Managed + @Flatten + public QueryManagerStats getStats() + { + return stats; + } + + public QueryId createQueryId() + { + return queryIdGenerator.createNextQueryId(); + } + + public ListenableFuture createQuery(QueryId queryId, String slug, SessionContext sessionContext, String query) + { + requireNonNull(queryId, "queryId is null"); + requireNonNull(sessionContext, "sessionFactory is null"); + requireNonNull(query, "query is null"); + checkArgument(!query.isEmpty(), "query must not be empty string"); + checkArgument(!queryTracker.tryGetQuery(queryId).isPresent(), "query %s already exists", queryId); + + DispatchQueryCreationFuture queryCreationFuture = new DispatchQueryCreationFuture(); + queryExecutor.execute(() -> { + try { + createQueryInternal(queryId, slug, sessionContext, query, resourceGroupManager); + } + finally { + queryCreationFuture.set(null); + } + }); + return queryCreationFuture; + } + + /** + * Creates and registers a dispatch query with the query tracker. This method will never fail to register a query with the query + * tracker. If an error occurs while, creating a dispatch query a failed dispatch will be created and registered. + */ + private void createQueryInternal(QueryId queryId, String slug, SessionContext sessionContext, String query, ResourceGroupManager resourceGroupManager) + { + Session session = null; + PreparedQuery preparedQuery; + try { + if (query.length() > maxQueryLength) { + int queryLength = query.length(); + query = query.substring(0, maxQueryLength); + throw new PrestoException(QUERY_TEXT_TOO_LARGE, format("Query text length (%s) exceeds the maximum length (%s)", queryLength, maxQueryLength)); + } + + // decode session + session = sessionSupplier.createSession(queryId, sessionContext); + + // prepare query + preparedQuery = queryPreparer.prepareQuery(session, query, warningCollector); + + // select resource group + Optional queryType = getQueryType(preparedQuery.getStatement().getClass()); + SelectionContext selectionContext = resourceGroupManager.selectGroup(new SelectionCriteria( + sessionContext.getIdentity().getPrincipal().isPresent(), + sessionContext.getIdentity().getUser(), + Optional.ofNullable(sessionContext.getSource()), + sessionContext.getClientTags(), + sessionContext.getResourceEstimates(), + queryType.map(Enum::name))); + + // apply system default session properties (does not override user set properties) + session = sessionPropertyDefaults.newSessionWithDefaultProperties(session, queryType.map(Enum::name), selectionContext.getResourceGroupId()); + + // mark existing transaction as active + transactionManager.activateTransaction(session, isTransactionControlStatement(preparedQuery.getStatement()), accessControl); + + DispatchQuery dispatchQuery = dispatchQueryFactory.createDispatchQuery( + session, + query, + preparedQuery, + slug, + selectionContext.getResourceGroupId(), + queryType, + warningCollector); + + boolean queryAdded = queryCreated(dispatchQuery); + if (queryAdded && !dispatchQuery.isDone()) { + try { + resourceGroupManager.submit(preparedQuery.getStatement(), dispatchQuery, selectionContext, queryExecutor); + } + catch (Throwable e) { + // dispatch query has already been registered, so just fail it directly + dispatchQuery.fail(e); + } + } + } + catch (Throwable throwable) { + // creation must never fail, so register a failed query in this case + if (session == null) { + session = Session.builder(new SessionPropertyManager()) + .setQueryId(queryId) + .setIdentity(sessionContext.getIdentity()) + .setSource(sessionContext.getSource()) + .build(); + } + DispatchQuery failedDispatchQuery = failedDispatchQueryFactory.createFailedDispatchQuery(session, query, Optional.empty(), throwable); + queryCreated(failedDispatchQuery); + } + } + + private boolean queryCreated(DispatchQuery dispatchQuery) + { + boolean queryAdded = queryTracker.addQuery(dispatchQuery); + + // only add state tracking if this query instance will actually be used for the execution + if (queryAdded) { + dispatchQuery.addStateChangeListener(newState -> { + if (newState.isDone()) { + // execution MUST be added to the expiration queue or there will be a leak + queryTracker.expireQuery(dispatchQuery.getQueryId()); + } + }); + stats.trackQueryStats(dispatchQuery); + } + + return queryAdded; + } + + public ListenableFuture waitForDispatched(QueryId queryId) + { + return queryTracker.tryGetQuery(queryId) + .map(dispatchQuery -> { + dispatchQuery.recordHeartbeat(); + return dispatchQuery.getDispatchedFuture(); + }) + .orElseGet(() -> immediateFuture(null)); + } + + public List getQueries() + { + return queryTracker.getAllQueries().stream() + .map(DispatchQuery::getBasicQueryInfo) + .collect(toImmutableList()); + } + + public BasicQueryInfo getQueryInfo(QueryId queryId) + { + return queryTracker.getQuery(queryId).getBasicQueryInfo(); + } + + public Optional getDispatchInfo(QueryId queryId) + { + return queryTracker.tryGetQuery(queryId) + .map(dispatchQuery -> { + dispatchQuery.recordHeartbeat(); + return dispatchQuery.getDispatchInfo(); + }); + } + + public void cancelQuery(QueryId queryId) + { + queryTracker.tryGetQuery(queryId) + .ifPresent(DispatchQuery::cancel); + } + + private static class DispatchQueryCreationFuture + extends AbstractFuture + { + @Override + protected boolean set(QueryInfo value) + { + return super.set(value); + } + + @Override + protected boolean setException(Throwable throwable) + { + return super.setException(throwable); + } + + @Override + public boolean cancel(boolean mayInterruptIfRunning) + { + // query submission can not be canceled + return false; + } + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/dispatcher/DispatchQuery.java b/presto-main/src/main/java/com/facebook/presto/dispatcher/DispatchQuery.java new file mode 100644 index 0000000000000..cedc15a2c12ea --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/dispatcher/DispatchQuery.java @@ -0,0 +1,30 @@ +/* + * 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 com.facebook.presto.dispatcher; + +import com.facebook.presto.execution.ManagedQueryExecution; +import com.facebook.presto.execution.QueryTracker.TrackedQuery; +import com.google.common.util.concurrent.ListenableFuture; + +public interface DispatchQuery + extends TrackedQuery, ManagedQueryExecution +{ + void recordHeartbeat(); + + ListenableFuture getDispatchedFuture(); + + DispatchInfo getDispatchInfo(); + + void cancel(); +} diff --git a/presto-main/src/main/java/com/facebook/presto/dispatcher/DispatchQueryFactory.java b/presto-main/src/main/java/com/facebook/presto/dispatcher/DispatchQueryFactory.java new file mode 100644 index 0000000000000..41253e64c4416 --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/dispatcher/DispatchQueryFactory.java @@ -0,0 +1,34 @@ +/* + * 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 com.facebook.presto.dispatcher; + +import com.facebook.presto.Session; +import com.facebook.presto.execution.QueryPreparer.PreparedQuery; +import com.facebook.presto.execution.warnings.WarningCollector; +import com.facebook.presto.spi.resourceGroups.QueryType; +import com.facebook.presto.spi.resourceGroups.ResourceGroupId; + +import java.util.Optional; + +public interface DispatchQueryFactory +{ + DispatchQuery createDispatchQuery( + Session session, + String query, + PreparedQuery preparedQuery, + String slug, + ResourceGroupId resourceGroup, + Optional queryType, + WarningCollector warningCollector); +} diff --git a/presto-main/src/main/java/com/facebook/presto/execution/FailedQueryExecution.java b/presto-main/src/main/java/com/facebook/presto/dispatcher/FailedDispatchQuery.java similarity index 54% rename from presto-main/src/main/java/com/facebook/presto/execution/FailedQueryExecution.java rename to presto-main/src/main/java/com/facebook/presto/dispatcher/FailedDispatchQuery.java index 417ca1578a0e9..8aa670a3562cd 100644 --- a/presto-main/src/main/java/com/facebook/presto/execution/FailedQueryExecution.java +++ b/presto-main/src/main/java/com/facebook/presto/dispatcher/FailedDispatchQuery.java @@ -11,17 +11,16 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package com.facebook.presto.execution; +package com.facebook.presto.dispatcher; import com.facebook.presto.Session; +import com.facebook.presto.execution.ExecutionFailureInfo; +import com.facebook.presto.execution.QueryState; import com.facebook.presto.execution.StateMachine.StateChangeListener; -import com.facebook.presto.memory.VersionedMemoryPoolId; import com.facebook.presto.server.BasicQueryInfo; import com.facebook.presto.spi.ErrorCode; import com.facebook.presto.spi.QueryId; -import com.facebook.presto.spi.resourceGroups.QueryType; import com.facebook.presto.spi.resourceGroups.ResourceGroupId; -import com.facebook.presto.sql.planner.Plan; import com.google.common.util.concurrent.ListenableFuture; import io.airlift.units.DataSize; import io.airlift.units.Duration; @@ -30,190 +29,149 @@ import java.net.URI; import java.util.Optional; import java.util.concurrent.Executor; -import java.util.function.Consumer; -import static com.facebook.presto.execution.QueryInfo.immediateFailureQueryInfo; import static com.facebook.presto.execution.QueryState.FAILED; -import static com.facebook.presto.memory.LocalMemoryManager.GENERAL_POOL; +import static com.facebook.presto.server.BasicQueryInfo.immediateFailureQueryInfo; import static com.google.common.util.concurrent.Futures.immediateFuture; import static io.airlift.units.DataSize.Unit.BYTE; import static java.util.Objects.requireNonNull; -import static java.util.concurrent.TimeUnit.NANOSECONDS; +import static java.util.concurrent.TimeUnit.MILLISECONDS; -public class FailedQueryExecution - implements QueryExecution +public class FailedDispatchQuery + implements DispatchQuery { - private final QueryInfo queryInfo; + private final BasicQueryInfo queryInfo; private final Session session; private final Executor executor; - - public FailedQueryExecution(Session session, String query, URI self, Optional resourceGroup, Optional queryType, Executor executor, Throwable cause) - { - requireNonNull(cause, "cause is null"); + private final DispatchInfo dispatchInfo; + + public FailedDispatchQuery( + Session session, + String query, + URI self, + Optional resourceGroup, + ExecutionFailureInfo failure, + Executor executor) + { + requireNonNull(session, "session is null"); + requireNonNull(query, "query is null"); + requireNonNull(self, "self is null"); + requireNonNull(resourceGroup, "resourceGroup is null"); + requireNonNull(failure, "failure is null"); + requireNonNull(executor, "executor is null"); + + this.queryInfo = immediateFailureQueryInfo(session, query, self, resourceGroup, failure.getErrorCode()); this.session = requireNonNull(session, "session is null"); this.executor = requireNonNull(executor, "executor is null"); - this.queryInfo = immediateFailureQueryInfo(session, query, self, resourceGroup, queryType, cause); - } - @Override - public QueryId getQueryId() - { - return queryInfo.getQueryId(); + this.dispatchInfo = DispatchInfo.failed( + failure, + queryInfo.getQueryStats().getElapsedTime(), + queryInfo.getQueryStats().getQueuedTime()); } @Override - public QueryInfo getQueryInfo() + public BasicQueryInfo getBasicQueryInfo() { return queryInfo; } @Override - public QueryState getState() - { - return queryInfo.getState(); - } - - @Override - public Plan getQueryPlan() - { - throw new UnsupportedOperationException(); - } - - @Override - public VersionedMemoryPoolId getMemoryPool() + public Session getSession() { - return new VersionedMemoryPoolId(GENERAL_POOL, 0); + return session; } @Override - public void setMemoryPool(VersionedMemoryPoolId poolId) + public ListenableFuture getDispatchedFuture() { - // no-op + return immediateFuture(null); } @Override - public DataSize getUserMemoryReservation() + public DispatchInfo getDispatchInfo() { - return new DataSize(0, BYTE); + return dispatchInfo; } @Override - public DataSize getTotalMemoryReservation() + public void addStateChangeListener(StateChangeListener stateChangeListener) { - return new DataSize(0, BYTE); + executor.execute(() -> stateChangeListener.stateChanged(FAILED)); } @Override - public Duration getTotalCpuTime() - { - return new Duration(0, NANOSECONDS); - } + public void startWaitingForResources() {} @Override - public Session getSession() - { - return session; - } + public void fail(Throwable throwable) {} @Override - public DateTime getCreateTime() - { - return queryInfo.getQueryStats().getCreateTime(); - } + public void cancel() {} @Override - public Optional getExecutionStartTime() - { - return Optional.ofNullable(queryInfo.getQueryStats().getExecutionStartTime()); - } + public void pruneInfo() {} @Override - public DateTime getLastHeartbeat() + public QueryId getQueryId() { - return queryInfo.getQueryStats().getLastHeartbeat(); + return queryInfo.getQueryId(); } @Override - public Optional getEndTime() + public boolean isDone() { - return Optional.ofNullable(queryInfo.getQueryStats().getEndTime()); + return true; } @Override public Optional getErrorCode() { - return Optional.ofNullable(getQueryInfo().getFailureInfo()).map(ExecutionFailureInfo::getErrorCode); - } - - @Override - public BasicQueryInfo getBasicQueryInfo() - { - return new BasicQueryInfo(getQueryInfo()); - } - - @Override - public void start() - { - // no-op - } - - @Override - public void addOutputInfoListener(Consumer listener) - { - // no-op - } - - @Override - public ListenableFuture getStateChange(QueryState currentState) - { - return immediateFuture(queryInfo.getState()); + return Optional.ofNullable(queryInfo.getErrorCode()); } @Override - public void addStateChangeListener(StateChangeListener stateChangeListener) - { - executor.execute(() -> stateChangeListener.stateChanged(FAILED)); - } + public void recordHeartbeat() {} @Override - public void addFinalQueryInfoListener(StateChangeListener stateChangeListener) + public DateTime getLastHeartbeat() { - executor.execute(() -> stateChangeListener.stateChanged(queryInfo)); + return queryInfo.getQueryStats().getEndTime(); } @Override - public void fail(Throwable cause) + public DateTime getCreateTime() { - // no-op + return queryInfo.getQueryStats().getCreateTime(); } @Override - public boolean isDone() + public Optional getExecutionStartTime() { - return getState().isDone(); + return getEndTime(); } @Override - public void cancelQuery() + public Optional getEndTime() { - // no-op + return Optional.ofNullable(queryInfo.getQueryStats().getEndTime()); } @Override - public void cancelStage(StageId stageId) + public Duration getTotalCpuTime() { - // no-op + return new Duration(0, MILLISECONDS); } @Override - public void recordHeartbeat() + public DataSize getTotalMemoryReservation() { - // no-op + return new DataSize(0, BYTE); } @Override - public void pruneInfo() + public DataSize getUserMemoryReservation() { - // no-op + return new DataSize(0, BYTE); } } diff --git a/presto-main/src/main/java/com/facebook/presto/dispatcher/FailedDispatchQueryFactory.java b/presto-main/src/main/java/com/facebook/presto/dispatcher/FailedDispatchQueryFactory.java new file mode 100644 index 0000000000000..173f7f6757e74 --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/dispatcher/FailedDispatchQueryFactory.java @@ -0,0 +1,63 @@ +/* + * 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 com.facebook.presto.dispatcher; + +import com.facebook.presto.Session; +import com.facebook.presto.event.QueryMonitor; +import com.facebook.presto.execution.ExecutionFailureInfo; +import com.facebook.presto.execution.LocationFactory; +import com.facebook.presto.server.BasicQueryInfo; +import com.facebook.presto.spi.resourceGroups.ResourceGroupId; + +import javax.inject.Inject; + +import java.util.Optional; +import java.util.concurrent.ExecutorService; + +import static com.facebook.presto.util.Failures.toFailure; +import static java.util.Objects.requireNonNull; + +public class FailedDispatchQueryFactory +{ + private final QueryMonitor queryMonitor; + private final LocationFactory locationFactory; + private final ExecutorService executor; + + @Inject + public FailedDispatchQueryFactory(QueryMonitor queryMonitor, LocationFactory locationFactory, DispatchExecutor dispatchExecutor) + { + this.queryMonitor = requireNonNull(queryMonitor, "queryMonitor is null"); + this.locationFactory = requireNonNull(locationFactory, "locationFactory is null"); + this.executor = requireNonNull(dispatchExecutor, "dispatchExecutor is null").getExecutor(); + } + + public FailedDispatchQuery createFailedDispatchQuery(Session session, String query, Optional resourceGroup, Throwable throwable) + { + ExecutionFailureInfo failure = toFailure(throwable); + FailedDispatchQuery failedDispatchQuery = new FailedDispatchQuery( + session, + query, + locationFactory.createQueryLocation(session.getQueryId()), + resourceGroup, + failure, + executor); + + BasicQueryInfo queryInfo = failedDispatchQuery.getBasicQueryInfo(); + + queryMonitor.queryCreatedEvent(queryInfo); + queryMonitor.queryImmediateFailureEvent(queryInfo, failure); + + return failedDispatchQuery; + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/dispatcher/LocalCoordinatorLocation.java b/presto-main/src/main/java/com/facebook/presto/dispatcher/LocalCoordinatorLocation.java new file mode 100644 index 0000000000000..badc3c34bd41f --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/dispatcher/LocalCoordinatorLocation.java @@ -0,0 +1,35 @@ +/* + * 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 com.facebook.presto.dispatcher; + +import javax.ws.rs.core.UriInfo; + +import java.net.URI; + +import static com.google.common.base.Strings.isNullOrEmpty; + +public class LocalCoordinatorLocation + implements CoordinatorLocation +{ + @Override + public URI getUri(UriInfo uriInfo, String xForwardedProto) + { + String scheme = isNullOrEmpty(xForwardedProto) ? uriInfo.getRequestUri().getScheme() : xForwardedProto; + return uriInfo.getRequestUriBuilder() + .scheme(scheme) + .replacePath("") + .replaceQuery("") + .build(); + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/dispatcher/LocalDispatchQuery.java b/presto-main/src/main/java/com/facebook/presto/dispatcher/LocalDispatchQuery.java new file mode 100644 index 0000000000000..1032fd59bd049 --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/dispatcher/LocalDispatchQuery.java @@ -0,0 +1,263 @@ +/* + * 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 com.facebook.presto.dispatcher; + +import com.facebook.presto.Session; +import com.facebook.presto.execution.ClusterSizeMonitor; +import com.facebook.presto.execution.ExecutionFailureInfo; +import com.facebook.presto.execution.QueryExecution; +import com.facebook.presto.execution.QueryState; +import com.facebook.presto.execution.QueryStateMachine; +import com.facebook.presto.execution.StateMachine.StateChangeListener; +import com.facebook.presto.server.BasicQueryInfo; +import com.facebook.presto.spi.ErrorCode; +import com.facebook.presto.spi.PrestoException; +import com.facebook.presto.spi.QueryId; +import com.google.common.util.concurrent.ListenableFuture; +import com.google.common.util.concurrent.SettableFuture; +import io.airlift.log.Logger; +import io.airlift.units.DataSize; +import io.airlift.units.Duration; +import org.joda.time.DateTime; + +import java.util.Optional; +import java.util.concurrent.Executor; +import java.util.function.Consumer; + +import static com.facebook.presto.execution.QueryState.FAILED; +import static com.facebook.presto.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR; +import static com.facebook.presto.util.Failures.toFailure; +import static com.google.common.util.concurrent.Futures.nonCancellationPropagating; +import static io.airlift.concurrent.MoreFutures.addExceptionCallback; +import static io.airlift.concurrent.MoreFutures.addSuccessCallback; +import static io.airlift.concurrent.MoreFutures.tryGetFutureValue; +import static io.airlift.units.DataSize.Unit.BYTE; +import static java.util.Objects.requireNonNull; +import static java.util.concurrent.TimeUnit.MILLISECONDS; + +public class LocalDispatchQuery + implements DispatchQuery +{ + private static final Logger log = Logger.get(LocalDispatchQuery.class); + private final QueryStateMachine stateMachine; + private final ListenableFuture queryExecutionFuture; + + private final ClusterSizeMonitor clusterSizeMonitor; + + private final Executor queryExecutor; + + private final Consumer querySubmitter; + private final SettableFuture submitted = SettableFuture.create(); + + public LocalDispatchQuery( + QueryStateMachine stateMachine, + ListenableFuture queryExecutionFuture, + ClusterSizeMonitor clusterSizeMonitor, + Executor queryExecutor, + Consumer querySubmitter) + { + this.stateMachine = requireNonNull(stateMachine, "stateMachine is null"); + this.queryExecutionFuture = requireNonNull(queryExecutionFuture, "queryExecutionFuture is null"); + this.clusterSizeMonitor = requireNonNull(clusterSizeMonitor, "clusterSizeMonitor is null"); + this.queryExecutor = requireNonNull(queryExecutor, "queryExecutor is null"); + this.querySubmitter = requireNonNull(querySubmitter, "querySubmitter is null"); + + addExceptionCallback(queryExecutionFuture, stateMachine::transitionToFailed); + stateMachine.addStateChangeListener(state -> { + if (state.isDone()) { + submitted.set(null); + } + }); + } + + @Override + public void startWaitingForResources() + { + if (stateMachine.transitionToWaitingForResources()) { + waitForMinimumWorkers(); + } + } + + private void waitForMinimumWorkers() + { + ListenableFuture minimumWorkerFuture = clusterSizeMonitor.waitForMinimumWorkers(); + // when worker requirement is met, wait for query execution to finish construction and then start the execution + addSuccessCallback(minimumWorkerFuture, () -> addSuccessCallback(queryExecutionFuture, this::startExecution)); + addExceptionCallback(minimumWorkerFuture, throwable -> queryExecutor.execute(() -> stateMachine.transitionToFailed(throwable))); + } + + private void startExecution(QueryExecution queryExecution) + { + queryExecutor.execute(() -> { + if (stateMachine.transitionToDispatching()) { + try { + querySubmitter.accept(queryExecution); + } + catch (Throwable t) { + // this should never happen but be safe + stateMachine.transitionToFailed(t); + log.error(t, "query submitter threw exception"); + throw t; + } + finally { + submitted.set(null); + } + } + }); + } + + @Override + public void recordHeartbeat() + { + stateMachine.recordHeartbeat(); + } + + @Override + public DateTime getLastHeartbeat() + { + return stateMachine.getLastHeartbeat(); + } + + @Override + public ListenableFuture getDispatchedFuture() + { + return nonCancellationPropagating(submitted); + } + + @Override + public DispatchInfo getDispatchInfo() + { + // observe submitted before getting the state, to ensure a failed query stat is visible + boolean dispatched = submitted.isDone(); + BasicQueryInfo queryInfo = stateMachine.getBasicQueryInfo(Optional.empty()); + + if (queryInfo.getState() == FAILED) { + ExecutionFailureInfo failureInfo = stateMachine.getFailureInfo() + .orElseGet(() -> toFailure(new PrestoException(GENERIC_INTERNAL_ERROR, "Query failed for an unknown reason"))); + return DispatchInfo.failed(failureInfo, queryInfo.getQueryStats().getElapsedTime(), queryInfo.getQueryStats().getQueuedTime()); + } + if (dispatched) { + return DispatchInfo.dispatched(new LocalCoordinatorLocation(), queryInfo.getQueryStats().getElapsedTime(), queryInfo.getQueryStats().getQueuedTime()); + } + return DispatchInfo.queued(queryInfo.getQueryStats().getElapsedTime(), queryInfo.getQueryStats().getQueuedTime()); + } + + @Override + public QueryId getQueryId() + { + return stateMachine.getQueryId(); + } + + @Override + public boolean isDone() + { + return stateMachine.getQueryState().isDone(); + } + + @Override + public DateTime getCreateTime() + { + return stateMachine.getCreateTime(); + } + + @Override + public Optional getExecutionStartTime() + { + return stateMachine.getExecutionStartTime(); + } + + @Override + public Optional getEndTime() + { + return stateMachine.getEndTime(); + } + + @Override + public Duration getTotalCpuTime() + { + return tryGetQueryExecution() + .map(QueryExecution::getTotalCpuTime) + .orElse(new Duration(0, MILLISECONDS)); + } + + @Override + public DataSize getTotalMemoryReservation() + { + return tryGetQueryExecution() + .map(QueryExecution::getTotalMemoryReservation) + .orElse(new DataSize(0, BYTE)); + } + + @Override + public DataSize getUserMemoryReservation() + { + return tryGetQueryExecution() + .map(QueryExecution::getUserMemoryReservation) + .orElse(new DataSize(0, BYTE)); + } + + @Override + public BasicQueryInfo getBasicQueryInfo() + { + return tryGetQueryExecution() + .map(QueryExecution::getBasicQueryInfo) + .orElse(stateMachine.getBasicQueryInfo(Optional.empty())); + } + + @Override + public Session getSession() + { + return stateMachine.getSession(); + } + + @Override + public void fail(Throwable throwable) + { + stateMachine.transitionToFailed(throwable); + } + + @Override + public void cancel() + { + stateMachine.transitionToCanceled(); + } + + @Override + public void pruneInfo() + { + stateMachine.pruneQueryInfo(); + } + + @Override + public Optional getErrorCode() + { + return stateMachine.getFailureInfo().map(ExecutionFailureInfo::getErrorCode); + } + + @Override + public void addStateChangeListener(StateChangeListener stateChangeListener) + { + stateMachine.addStateChangeListener(stateChangeListener); + } + + private Optional tryGetQueryExecution() + { + try { + return tryGetFutureValue(queryExecutionFuture); + } + catch (Exception ignored) { + return Optional.empty(); + } + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/dispatcher/LocalDispatchQueryFactory.java b/presto-main/src/main/java/com/facebook/presto/dispatcher/LocalDispatchQueryFactory.java new file mode 100644 index 0000000000000..69f05e68bf675 --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/dispatcher/LocalDispatchQueryFactory.java @@ -0,0 +1,126 @@ +/* + * 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 com.facebook.presto.dispatcher; + +import com.facebook.presto.Session; +import com.facebook.presto.event.QueryMonitor; +import com.facebook.presto.execution.ClusterSizeMonitor; +import com.facebook.presto.execution.LocationFactory; +import com.facebook.presto.execution.QueryExecution; +import com.facebook.presto.execution.QueryExecution.QueryExecutionFactory; +import com.facebook.presto.execution.QueryManager; +import com.facebook.presto.execution.QueryPreparer.PreparedQuery; +import com.facebook.presto.execution.QueryStateMachine; +import com.facebook.presto.execution.warnings.WarningCollector; +import com.facebook.presto.metadata.Metadata; +import com.facebook.presto.security.AccessControl; +import com.facebook.presto.spi.PrestoException; +import com.facebook.presto.spi.resourceGroups.QueryType; +import com.facebook.presto.spi.resourceGroups.ResourceGroupId; +import com.facebook.presto.sql.tree.Statement; +import com.facebook.presto.transaction.TransactionManager; +import com.google.common.util.concurrent.ListenableFuture; +import com.google.common.util.concurrent.ListeningExecutorService; + +import javax.inject.Inject; + +import java.util.Map; +import java.util.Optional; + +import static com.facebook.presto.spi.StandardErrorCode.NOT_SUPPORTED; +import static com.facebook.presto.util.StatementUtils.isTransactionControlStatement; +import static java.util.Objects.requireNonNull; + +public class LocalDispatchQueryFactory + implements DispatchQueryFactory +{ + private final QueryManager queryManager; + private final TransactionManager transactionManager; + private final AccessControl accessControl; + private final Metadata metadata; + private final QueryMonitor queryMonitor; + private final LocationFactory locationFactory; + + private final ClusterSizeMonitor clusterSizeMonitor; + + private final Map, QueryExecutionFactory> executionFactories; + private final ListeningExecutorService executor; + + @Inject + public LocalDispatchQueryFactory( + QueryManager queryManager, + TransactionManager transactionManager, + AccessControl accessControl, + Metadata metadata, + QueryMonitor queryMonitor, + LocationFactory locationFactory, + Map, QueryExecutionFactory> executionFactories, + ClusterSizeMonitor clusterSizeMonitor, + DispatchExecutor dispatchExecutor) + { + this.queryManager = requireNonNull(queryManager, "queryManager is null"); + this.transactionManager = requireNonNull(transactionManager, "transactionManager is null"); + this.accessControl = requireNonNull(accessControl, "accessControl is null"); + this.metadata = requireNonNull(metadata, "metadata is null"); + this.queryMonitor = requireNonNull(queryMonitor, "queryMonitor is null"); + this.locationFactory = requireNonNull(locationFactory, "locationFactory is null"); + this.executionFactories = requireNonNull(executionFactories, "executionFactories is null"); + + this.clusterSizeMonitor = requireNonNull(clusterSizeMonitor, "clusterSizeMonitor is null"); + + this.executor = requireNonNull(dispatchExecutor, "executorService is null").getExecutor(); + } + + @Override + public DispatchQuery createDispatchQuery( + Session session, + String query, + PreparedQuery preparedQuery, + String slug, + ResourceGroupId resourceGroup, + Optional queryType, + WarningCollector warningCollector) + { + QueryStateMachine stateMachine = QueryStateMachine.begin( + query, + session, + locationFactory.createQueryLocation(session.getQueryId()), + resourceGroup, + queryType, + isTransactionControlStatement(preparedQuery.getStatement()), + transactionManager, + accessControl, + executor, + metadata, + warningCollector); + + queryMonitor.queryCreatedEvent(stateMachine.getBasicQueryInfo(Optional.empty())); + + ListenableFuture queryExecutionFuture = executor.submit(() -> { + QueryExecutionFactory queryExecutionFactory = executionFactories.get(preparedQuery.getStatement().getClass()); + if (queryExecutionFactory == null) { + throw new PrestoException(NOT_SUPPORTED, "Unsupported statement type: " + preparedQuery.getStatement().getClass().getSimpleName()); + } + + return queryExecutionFactory.createQueryExecution(preparedQuery, stateMachine, slug, warningCollector, queryType); + }); + + return new LocalDispatchQuery( + stateMachine, + queryExecutionFuture, + clusterSizeMonitor, + executor, + queryManager::createQuery); + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/dispatcher/QueuedStatementResource.java b/presto-main/src/main/java/com/facebook/presto/dispatcher/QueuedStatementResource.java new file mode 100644 index 0000000000000..e95ed28606269 --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/dispatcher/QueuedStatementResource.java @@ -0,0 +1,436 @@ +/* + * 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 com.facebook.presto.dispatcher; + +import com.facebook.presto.client.QueryError; +import com.facebook.presto.client.QueryResults; +import com.facebook.presto.client.StatementStats; +import com.facebook.presto.execution.ExecutionFailureInfo; +import com.facebook.presto.execution.QueryState; +import com.facebook.presto.server.HttpRequestSessionContext; +import com.facebook.presto.server.SessionContext; +import com.facebook.presto.spi.ErrorCode; +import com.facebook.presto.spi.QueryId; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import com.google.common.collect.Ordering; +import com.google.common.util.concurrent.Futures; +import com.google.common.util.concurrent.ListenableFuture; +import io.airlift.log.Logger; +import io.airlift.units.Duration; + +import javax.annotation.PreDestroy; +import javax.annotation.concurrent.GuardedBy; +import javax.inject.Inject; +import javax.servlet.http.HttpServletRequest; +import javax.ws.rs.DELETE; +import javax.ws.rs.GET; +import javax.ws.rs.HeaderParam; +import javax.ws.rs.POST; +import javax.ws.rs.Path; +import javax.ws.rs.PathParam; +import javax.ws.rs.Produces; +import javax.ws.rs.QueryParam; +import javax.ws.rs.WebApplicationException; +import javax.ws.rs.container.AsyncResponse; +import javax.ws.rs.container.Suspended; +import javax.ws.rs.core.Context; +import javax.ws.rs.core.Response; +import javax.ws.rs.core.Response.Status; +import javax.ws.rs.core.UriInfo; + +import java.net.URI; +import java.util.Map.Entry; +import java.util.Optional; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.Executor; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.atomic.AtomicLong; + +import static com.facebook.presto.execution.QueryState.FAILED; +import static com.facebook.presto.execution.QueryState.QUEUED; +import static com.facebook.presto.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR; +import static com.google.common.base.MoreObjects.firstNonNull; +import static com.google.common.base.Strings.isNullOrEmpty; +import static com.google.common.net.HttpHeaders.X_FORWARDED_PROTO; +import static com.google.common.util.concurrent.MoreExecutors.directExecutor; +import static io.airlift.concurrent.MoreFutures.addTimeout; +import static io.airlift.concurrent.Threads.threadsNamed; +import static io.airlift.http.client.HttpUriBuilder.uriBuilderFrom; +import static io.airlift.http.server.AsyncResponseHandler.bindAsyncResponse; +import static java.util.Locale.ENGLISH; +import static java.util.Objects.requireNonNull; +import static java.util.UUID.randomUUID; +import static java.util.concurrent.Executors.newSingleThreadScheduledExecutor; +import static java.util.concurrent.TimeUnit.MILLISECONDS; +import static java.util.concurrent.TimeUnit.SECONDS; +import static javax.ws.rs.core.MediaType.APPLICATION_JSON; +import static javax.ws.rs.core.MediaType.TEXT_PLAIN_TYPE; +import static javax.ws.rs.core.Response.Status.BAD_REQUEST; +import static javax.ws.rs.core.Response.Status.NOT_FOUND; + +@Path("/") +public class QueuedStatementResource +{ + private static final Logger log = Logger.get(QueuedStatementResource.class); + private static final Duration MAX_WAIT_TIME = new Duration(1, SECONDS); + private static final Ordering> WAIT_ORDERING = Ordering.natural().nullsLast(); + private static final Duration NO_DURATION = new Duration(0, MILLISECONDS); + + private final DispatchManager dispatchManager; + + private final Executor responseExecutor; + private final ScheduledExecutorService timeoutExecutor; + + private final ConcurrentMap queries = new ConcurrentHashMap<>(); + private final ScheduledExecutorService queryPurger = newSingleThreadScheduledExecutor(threadsNamed("query-purger")); + + @Inject + public QueuedStatementResource( + DispatchManager dispatchManager, + DispatchExecutor executor) + { + this.dispatchManager = requireNonNull(dispatchManager, "dispatchManager is null"); + + requireNonNull(dispatchManager, "dispatchManager is null"); + this.responseExecutor = requireNonNull(executor, "responseExecutor is null").getExecutor(); + this.timeoutExecutor = requireNonNull(executor, "timeoutExecutor is null").getScheduledExecutor(); + + queryPurger.scheduleWithFixedDelay( + () -> { + try { + // snapshot the queries before checking states to avoid registration race + for (Entry entry : ImmutableSet.copyOf(queries.entrySet())) { + if (!entry.getValue().isSubmissionFinished()) { + continue; + } + + // forget about this query if the query manager is no longer tracking it + if (!dispatchManager.getDispatchInfo(entry.getKey()).isPresent()) { + queries.remove(entry.getKey()); + } + } + } + catch (Throwable e) { + log.warn(e, "Error removing old queries"); + } + }, + 200, + 200, + MILLISECONDS); + } + + @PreDestroy + public void stop() + { + queryPurger.shutdownNow(); + } + + @POST + @Path("/v1/statement") + @Produces(APPLICATION_JSON) + public Response postStatement( + String statement, + @HeaderParam(X_FORWARDED_PROTO) String xForwardedProto, + @Context HttpServletRequest servletRequest, + @Context UriInfo uriInfo) + { + if (isNullOrEmpty(statement)) { + throw badRequest(BAD_REQUEST, "SQL statement is empty"); + } + + SessionContext sessionContext = new HttpRequestSessionContext(servletRequest); + Query query = new Query(statement, sessionContext, dispatchManager); + queries.put(query.getQueryId(), query); + + return Response.ok(query.getQueryResults(query.getLastToken(), uriInfo, xForwardedProto)).build(); + } + + @GET + @Path("/v1/statement/queued/{queryId}/{slug}/{token}") + @Produces(APPLICATION_JSON) + public void getStatus( + @PathParam("queryId") QueryId queryId, + @PathParam("slug") String slug, + @PathParam("token") long token, + @QueryParam("maxWait") Duration maxWait, + @HeaderParam(X_FORWARDED_PROTO) String xForwardedProto, + @Context UriInfo uriInfo, + @Suspended AsyncResponse asyncResponse) + { + Query query = getQuery(queryId, slug); + + // wait for query to be dispatched, up to the wait timeout + ListenableFuture futureStateChange = addTimeout( + query.waitForDispatched(), + () -> null, + WAIT_ORDERING.min(MAX_WAIT_TIME, maxWait), + timeoutExecutor); + + // when state changes, fetch the next result + ListenableFuture queryResultsFuture = Futures.transform( + futureStateChange, + ignored -> query.getQueryResults(token, uriInfo, xForwardedProto), + responseExecutor); + + // transform to Response + ListenableFuture response = Futures.transform( + queryResultsFuture, + queryResults -> Response.ok(queryResults).build(), + directExecutor()); + bindAsyncResponse(asyncResponse, response, responseExecutor); + } + + @DELETE + @Path("/v1/statement/queued/{queryId}/{slug}/{token}") + @Produces(APPLICATION_JSON) + public Response cancelQuery( + @PathParam("queryId") QueryId queryId, + @PathParam("slug") String slug, + @PathParam("token") long token) + { + getQuery(queryId, slug) + .cancel(); + return Response.noContent().build(); + } + + private Query getQuery(QueryId queryId, String slug) + { + Query query = queries.get(queryId); + if (query == null || !query.getSlug().equals(slug)) { + throw badRequest(NOT_FOUND, "Query not found"); + } + return query; + } + + private static URI getQueryHtmlUri(QueryId queryId, UriInfo uriInfo, String xForwardedProto) + { + return uriInfo.getRequestUriBuilder() + .scheme(getScheme(xForwardedProto, uriInfo)) + .replacePath("ui/query.html") + .replaceQuery(queryId.toString()) + .build(); + } + + private static URI getQueuedUri(QueryId queryId, String slug, long token, UriInfo uriInfo, String xForwardedProto) + { + return uriInfo.getBaseUriBuilder() + .scheme(getScheme(xForwardedProto, uriInfo)) + .replacePath("/v1/statement/queued/") + .path(queryId.toString()) + .path(slug) + .path(String.valueOf(token)) + .replaceQuery("") + .build(); + } + + private static String getScheme(String xForwardedProto, @Context UriInfo uriInfo) + { + return isNullOrEmpty(xForwardedProto) ? uriInfo.getRequestUri().getScheme() : xForwardedProto; + } + + private static QueryResults createQueryResults( + QueryId queryId, + URI nextUri, + Optional queryError, + UriInfo uriInfo, + String xForwardedProto, + Duration elapsedTime, + Duration queuedTime) + { + QueryState state = queryError.map(error -> FAILED).orElse(QUEUED); + return new QueryResults( + queryId.toString(), + getQueryHtmlUri(queryId, uriInfo, xForwardedProto), + null, + nextUri, + null, + null, + StatementStats.builder() + .setState(state.toString()) + .setQueued(state == QUEUED) + .setElapsedTimeMillis(elapsedTime.toMillis()) + .setQueuedTimeMillis(queuedTime.toMillis()) + .build(), + queryError.orElse(null), + ImmutableList.of(), + null, + null); + } + + private static WebApplicationException badRequest(Status status, String message) + { + throw new WebApplicationException( + Response.status(status) + .type(TEXT_PLAIN_TYPE) + .entity(message) + .build()); + } + + private static final class Query + { + private final String query; + private final SessionContext sessionContext; + private final DispatchManager dispatchManager; + private final QueryId queryId; + private final String slug = "x" + randomUUID().toString().toLowerCase(ENGLISH).replace("-", ""); + private final AtomicLong lastToken = new AtomicLong(); + + @GuardedBy("this") + private ListenableFuture querySubmissionFuture; + + public Query(String query, SessionContext sessionContext, DispatchManager dispatchManager) + { + this.query = requireNonNull(query, "query is null"); + this.sessionContext = requireNonNull(sessionContext, "sessionContext is null"); + this.dispatchManager = requireNonNull(dispatchManager, "dispatchManager is null"); + this.queryId = dispatchManager.createQueryId(); + } + + public QueryId getQueryId() + { + return queryId; + } + + public String getSlug() + { + return slug; + } + + public long getLastToken() + { + return lastToken.get(); + } + + public synchronized boolean isSubmissionFinished() + { + return querySubmissionFuture != null && querySubmissionFuture.isDone(); + } + + private ListenableFuture waitForDispatched() + { + // if query query submission has not finished, wait for it to finish + synchronized (this) { + if (querySubmissionFuture == null) { + querySubmissionFuture = dispatchManager.createQuery(queryId, slug, sessionContext, query); + } + if (!querySubmissionFuture.isDone()) { + return querySubmissionFuture; + } + } + + // otherwise, wait for the query to finish + return dispatchManager.waitForDispatched(queryId); + } + + public QueryResults getQueryResults(long token, UriInfo uriInfo, String xForwardedProto) + { + long lastToken = this.lastToken.get(); + // token should be the last token or the next token + if (token != lastToken && token != lastToken + 1) { + throw new WebApplicationException(Response.Status.GONE); + } + // advance (or stay at) the token + this.lastToken.compareAndSet(lastToken, token); + + synchronized (this) { + // if query submission has not finished, return simple empty result + if (querySubmissionFuture == null || !querySubmissionFuture.isDone()) { + return createQueryResults( + token + 1, + uriInfo, + xForwardedProto, + DispatchInfo.queued(NO_DURATION, NO_DURATION)); + } + } + + Optional dispatchInfo = dispatchManager.getDispatchInfo(queryId); + if (!dispatchInfo.isPresent()) { + // query should always be found, but it may have just been determined to be abandoned + throw new WebApplicationException(Response + .status(NOT_FOUND) + .build()); + } + + return createQueryResults(token + 1, uriInfo, xForwardedProto, dispatchInfo.get()); + } + + public synchronized void cancel() + { + querySubmissionFuture.addListener(() -> dispatchManager.cancelQuery(queryId), directExecutor()); + } + + private QueryResults createQueryResults(long token, UriInfo uriInfo, String xForwardedProto, DispatchInfo dispatchInfo) + { + URI nextUri = getNextUri(token, uriInfo, xForwardedProto, dispatchInfo); + + Optional queryError = dispatchInfo.getFailureInfo() + .map(this::toQueryError); + + return QueuedStatementResource.createQueryResults( + queryId, + nextUri, + queryError, + uriInfo, + xForwardedProto, + dispatchInfo.getElapsedTime(), + dispatchInfo.getQueuedTime()); + } + + private URI getNextUri(long token, UriInfo uriInfo, String xForwardedProto, DispatchInfo dispatchInfo) + { + // if failed, query is complete + if (dispatchInfo.getFailureInfo().isPresent()) { + return null; + } + // if dispatched, redirect to new uri + return dispatchInfo.getCoordinatorLocation() + .map(coordinatorLocation -> getRedirectUri(coordinatorLocation, uriInfo, xForwardedProto)) + .orElseGet(() -> getQueuedUri(queryId, slug, token, uriInfo, xForwardedProto)); + } + + private URI getRedirectUri(CoordinatorLocation coordinatorLocation, UriInfo uriInfo, String xForwardedProto) + { + URI coordinatorUri = coordinatorLocation.getUri(uriInfo, xForwardedProto); + return uriBuilderFrom(coordinatorUri) + .appendPath("/v1/statement/executing") + .appendPath(queryId.toString()) + .appendPath(slug) + .appendPath("0") + .build(); + } + + private QueryError toQueryError(ExecutionFailureInfo executionFailureInfo) + { + ErrorCode errorCode; + if (executionFailureInfo.getErrorCode() != null) { + errorCode = executionFailureInfo.getErrorCode(); + } + else { + errorCode = GENERIC_INTERNAL_ERROR.toErrorCode(); + log.warn("Failed query %s has no error code", queryId); + } + + return new QueryError( + firstNonNull(executionFailureInfo.getMessage(), "Internal error"), + null, + errorCode.getCode(), + errorCode.getName(), + errorCode.getType().toString(), + executionFailureInfo.getErrorLocation(), + executionFailureInfo.toFailureInfo()); + } + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/event/QueryMonitor.java b/presto-main/src/main/java/com/facebook/presto/event/QueryMonitor.java index 2cb64dbd03815..f8d29b8f80003 100644 --- a/presto-main/src/main/java/com/facebook/presto/event/QueryMonitor.java +++ b/presto-main/src/main/java/com/facebook/presto/event/QueryMonitor.java @@ -148,7 +148,7 @@ public void queryImmediateFailureEvent(BasicQueryInfo queryInfo, ExecutionFailur new QueryStatistics( ofMillis(0), ofMillis(0), - ofMillis(0), + ofMillis(queryInfo.getQueryStats().getQueuedTime().toMillis()), Optional.empty(), 0, 0, @@ -173,7 +173,7 @@ public void queryImmediateFailureEvent(BasicQueryInfo queryInfo, ExecutionFailur ImmutableList.of(), queryInfo.getQueryType(), ofEpochMilli(queryInfo.getQueryStats().getCreateTime().getMillis()), - ofEpochMilli(queryInfo.getQueryStats().getCreateTime().getMillis()), + ofEpochMilli(queryInfo.getQueryStats().getEndTime().getMillis()), ofEpochMilli(queryInfo.getQueryStats().getEndTime().getMillis()))); logQueryTimeline(queryInfo); diff --git a/presto-main/src/main/java/com/facebook/presto/execution/ClusterSizeMonitor.java b/presto-main/src/main/java/com/facebook/presto/execution/ClusterSizeMonitor.java index bafe9512443e2..cec5e284ba847 100644 --- a/presto-main/src/main/java/com/facebook/presto/execution/ClusterSizeMonitor.java +++ b/presto-main/src/main/java/com/facebook/presto/execution/ClusterSizeMonitor.java @@ -35,11 +35,9 @@ import java.util.function.Consumer; import static com.facebook.presto.spi.StandardErrorCode.GENERIC_INSUFFICIENT_RESOURCES; -import static com.facebook.presto.spi.StandardErrorCode.SERVER_STARTING_UP; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.util.concurrent.Futures.immediateFuture; import static io.airlift.concurrent.Threads.threadsNamed; -import static io.airlift.units.Duration.nanosSince; import static java.lang.String.format; import static java.util.Objects.requireNonNull; import static java.util.concurrent.Executors.newSingleThreadScheduledExecutor; @@ -49,14 +47,10 @@ public class ClusterSizeMonitor { private final InternalNodeManager nodeManager; private final boolean includeCoordinator; - private final int initializationMinCount; - private final Duration initializationMaxWait; private final int executionMinCount; private final Duration executionMaxWait; private final ScheduledExecutorService executor; - private final long createNanos = System.nanoTime(); - private final Consumer listener = this::updateAllNodes; @GuardedBy("this") @@ -65,34 +59,24 @@ public class ClusterSizeMonitor @GuardedBy("this") private final List> futures = new ArrayList<>(); - @GuardedBy("this") - private boolean minimumWorkerRequirementMet; - @Inject public ClusterSizeMonitor(InternalNodeManager nodeManager, NodeSchedulerConfig nodeSchedulerConfig, QueryManagerConfig queryManagerConfig) { this( nodeManager, requireNonNull(nodeSchedulerConfig, "nodeSchedulerConfig is null").isIncludeCoordinator(), - requireNonNull(queryManagerConfig, "queryManagerConfig is null").getInitializationRequiredWorkers(), - queryManagerConfig.getInitializationTimeout(), - queryManagerConfig.getRequiredWorkers(), + requireNonNull(queryManagerConfig, "queryManagerConfig is null").getRequiredWorkers(), queryManagerConfig.getRequiredWorkersMaxWait()); } public ClusterSizeMonitor( InternalNodeManager nodeManager, boolean includeCoordinator, - int initializationMinCount, - Duration initializationMaxWait, int executionMinCount, Duration executionMaxWait) { this.nodeManager = requireNonNull(nodeManager, "nodeManager is null"); this.includeCoordinator = includeCoordinator; - checkArgument(initializationMinCount >= 0, "initializationMinCount is negative"); - this.initializationMinCount = initializationMinCount; - this.initializationMaxWait = requireNonNull(initializationMaxWait, "initializationMaxWait is null"); checkArgument(executionMinCount >= 0, "executionMinCount is negative"); this.executionMinCount = executionMinCount; this.executionMaxWait = requireNonNull(executionMaxWait, "executionMaxWait is null"); @@ -112,18 +96,6 @@ public void stop() nodeManager.removeNodeChangeListener(listener); } - public synchronized void verifyInitialMinimumWorkersRequirement() - { - if (minimumWorkerRequirementMet) { - return; - } - - if (currentCount < initializationMinCount && nanosSince(createNanos).compareTo(initializationMaxWait) < 0) { - throw new PrestoException(SERVER_STARTING_UP, format("Cluster is still initializing, there are insufficient active worker nodes (%s) to run query", currentCount)); - } - minimumWorkerRequirementMet = true; - } - /** * Returns a listener that completes when the minimum number of workers for the cluster has been met. * Note: caller should not add a listener using the direct executor, as this can delay the diff --git a/presto-main/src/main/java/com/facebook/presto/execution/DataDefinitionExecution.java b/presto-main/src/main/java/com/facebook/presto/execution/DataDefinitionExecution.java index 75b4e74cd871e..9817486425ef1 100644 --- a/presto-main/src/main/java/com/facebook/presto/execution/DataDefinitionExecution.java +++ b/presto-main/src/main/java/com/facebook/presto/execution/DataDefinitionExecution.java @@ -22,10 +22,8 @@ import com.facebook.presto.metadata.MetadataManager; import com.facebook.presto.security.AccessControl; import com.facebook.presto.server.BasicQueryInfo; -import com.facebook.presto.spi.ErrorCode; import com.facebook.presto.spi.QueryId; import com.facebook.presto.spi.resourceGroups.QueryType; -import com.facebook.presto.spi.resourceGroups.ResourceGroupId; import com.facebook.presto.sql.planner.Plan; import com.facebook.presto.sql.tree.Expression; import com.facebook.presto.sql.tree.Statement; @@ -43,7 +41,6 @@ import java.util.List; import java.util.Map; import java.util.Optional; -import java.util.concurrent.ExecutorService; import java.util.function.Consumer; import static com.google.common.base.Preconditions.checkArgument; @@ -58,6 +55,7 @@ public class DataDefinitionExecution { private final DataDefinitionTask task; private final T statement; + private final String slug; private final TransactionManager transactionManager; private final Metadata metadata; private final AccessControl accessControl; @@ -67,6 +65,7 @@ public class DataDefinitionExecution private DataDefinitionExecution( DataDefinitionTask task, T statement, + String slug, TransactionManager transactionManager, Metadata metadata, AccessControl accessControl, @@ -75,6 +74,7 @@ private DataDefinitionExecution( { this.task = requireNonNull(task, "task is null"); this.statement = requireNonNull(statement, "statement is null"); + this.slug = requireNonNull(slug, "slug is null"); this.transactionManager = requireNonNull(transactionManager, "transactionManager is null"); this.metadata = requireNonNull(metadata, "metadata is null"); this.accessControl = requireNonNull(accessControl, "accessControl is null"); @@ -82,6 +82,12 @@ private DataDefinitionExecution( this.parameters = parameters; } + @Override + public String getSlug() + { + return slug; + } + @Override public VersionedMemoryPoolId getMemoryPool() { @@ -100,12 +106,6 @@ public Session getSession() return stateMachine.getSession(); } - @Override - public Optional getErrorCode() - { - return stateMachine.getFailureInfo().map(ExecutionFailureInfo::getErrorCode); - } - @Override public DataSize getUserMemoryReservation() { @@ -280,69 +280,47 @@ public List getParameters() public static class DataDefinitionExecutionFactory implements QueryExecutionFactory> { - private final LocationFactory locationFactory; private final TransactionManager transactionManager; private final Metadata metadata; private final AccessControl accessControl; - private final ExecutorService executor; private final Map, DataDefinitionTask> tasks; @Inject public DataDefinitionExecutionFactory( - LocationFactory locationFactory, TransactionManager transactionManager, MetadataManager metadata, AccessControl accessControl, - @ForQueryExecution ExecutorService executor, Map, DataDefinitionTask> tasks) { - this.locationFactory = requireNonNull(locationFactory, "locationFactory is null"); this.transactionManager = requireNonNull(transactionManager, "transactionManager is null"); this.metadata = requireNonNull(metadata, "metadata is null"); this.accessControl = requireNonNull(accessControl, "accessControl is null"); - this.executor = requireNonNull(executor, "executor is null"); this.tasks = requireNonNull(tasks, "tasks is null"); } @Override public DataDefinitionExecution createQueryExecution( - String query, - Session session, PreparedQuery preparedQuery, - ResourceGroupId resourceGroup, + QueryStateMachine stateMachine, + String slug, WarningCollector warningCollector, Optional queryType) { - return createDataDefinitionExecution(query, session, resourceGroup, preparedQuery.getStatement(), preparedQuery.getParameters(), warningCollector, queryType); + return createDataDefinitionExecution(preparedQuery.getStatement(), preparedQuery.getParameters(), stateMachine, slug); } private DataDefinitionExecution createDataDefinitionExecution( - String query, - Session session, - ResourceGroupId resourceGroup, T statement, List parameters, - WarningCollector warningCollector, - Optional queryType) + QueryStateMachine stateMachine, + String slug) { @SuppressWarnings("unchecked") DataDefinitionTask task = (DataDefinitionTask) tasks.get(statement.getClass()); checkArgument(task != null, "no task for statement: %s", statement.getClass().getSimpleName()); - QueryStateMachine stateMachine = QueryStateMachine.begin( - query, - session, - locationFactory.createQueryLocation(session.getQueryId()), - resourceGroup, - queryType, - task.isTransactionControl(), - transactionManager, - accessControl, - executor, - metadata, - warningCollector); stateMachine.setUpdateType(task.getName()); - return new DataDefinitionExecution<>(task, statement, transactionManager, metadata, accessControl, stateMachine, parameters); + return new DataDefinitionExecution<>(task, statement, slug, transactionManager, metadata, accessControl, stateMachine, parameters); } } } diff --git a/presto-main/src/main/java/com/facebook/presto/execution/ManagedQueryExecution.java b/presto-main/src/main/java/com/facebook/presto/execution/ManagedQueryExecution.java index c116110e39c3c..b23b834c32dcb 100644 --- a/presto-main/src/main/java/com/facebook/presto/execution/ManagedQueryExecution.java +++ b/presto-main/src/main/java/com/facebook/presto/execution/ManagedQueryExecution.java @@ -24,7 +24,7 @@ public interface ManagedQueryExecution { - void start(); + void startWaitingForResources(); void fail(Throwable cause); diff --git a/presto-main/src/main/java/com/facebook/presto/execution/QueryExecution.java b/presto-main/src/main/java/com/facebook/presto/execution/QueryExecution.java index ad9369675a696..e8d26144c8290 100644 --- a/presto-main/src/main/java/com/facebook/presto/execution/QueryExecution.java +++ b/presto-main/src/main/java/com/facebook/presto/execution/QueryExecution.java @@ -13,19 +13,20 @@ */ package com.facebook.presto.execution; -import com.facebook.presto.Session; import com.facebook.presto.execution.QueryPreparer.PreparedQuery; import com.facebook.presto.execution.QueryTracker.TrackedQuery; import com.facebook.presto.execution.StateMachine.StateChangeListener; import com.facebook.presto.execution.warnings.WarningCollector; import com.facebook.presto.memory.VersionedMemoryPoolId; +import com.facebook.presto.server.BasicQueryInfo; import com.facebook.presto.spi.resourceGroups.QueryType; -import com.facebook.presto.spi.resourceGroups.ResourceGroupId; import com.facebook.presto.spi.type.Type; import com.facebook.presto.sql.planner.Plan; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.util.concurrent.ListenableFuture; +import io.airlift.units.DataSize; +import io.airlift.units.Duration; import java.net.URI; import java.util.List; @@ -36,22 +37,36 @@ import static java.util.Objects.requireNonNull; public interface QueryExecution - extends ManagedQueryExecution, TrackedQuery + extends TrackedQuery { QueryState getState(); ListenableFuture getStateChange(QueryState currentState); + void addStateChangeListener(StateChangeListener stateChangeListener); + void addOutputInfoListener(Consumer listener); Plan getQueryPlan(); + BasicQueryInfo getBasicQueryInfo(); + QueryInfo getQueryInfo(); + String getSlug(); + + Duration getTotalCpuTime(); + + DataSize getUserMemoryReservation(); + + DataSize getTotalMemoryReservation(); + VersionedMemoryPoolId getMemoryPool(); void setMemoryPool(VersionedMemoryPoolId poolId); + void start(); + void cancelQuery(); void cancelStage(StageId stageId); @@ -68,10 +83,9 @@ public interface QueryExecution interface QueryExecutionFactory { T createQueryExecution( - String query, - Session session, PreparedQuery preparedQuery, - ResourceGroupId resourceGroup, + QueryStateMachine stateMachine, + String slug, WarningCollector warningCollector, Optional queryType); } diff --git a/presto-main/src/main/java/com/facebook/presto/execution/QueryInfo.java b/presto-main/src/main/java/com/facebook/presto/execution/QueryInfo.java index 16c3c5fccb799..ead00da832d3c 100644 --- a/presto-main/src/main/java/com/facebook/presto/execution/QueryInfo.java +++ b/presto-main/src/main/java/com/facebook/presto/execution/QueryInfo.java @@ -13,7 +13,6 @@ */ package com.facebook.presto.execution; -import com.facebook.presto.Session; import com.facebook.presto.SessionRepresentation; import com.facebook.presto.spi.ErrorCode; import com.facebook.presto.spi.ErrorType; @@ -39,11 +38,7 @@ import java.util.Optional; import java.util.Set; -import static com.facebook.presto.execution.QueryState.FAILED; -import static com.facebook.presto.execution.QueryStats.immediateFailureQueryStats; import static com.facebook.presto.execution.StageInfo.getAllStages; -import static com.facebook.presto.memory.LocalMemoryManager.GENERAL_POOL; -import static com.facebook.presto.util.Failures.toFailure; import static com.google.common.base.MoreObjects.toStringHelper; import static java.util.Objects.requireNonNull; @@ -167,49 +162,6 @@ public QueryInfo( this.queryType = queryType; } - public static QueryInfo immediateFailureQueryInfo( - Session session, - String query, - URI self, - Optional resourceGroupId, - Optional queryType, - Throwable throwable) - { - ExecutionFailureInfo failureCause = toFailure(throwable); - QueryInfo queryInfo = new QueryInfo( - session.getQueryId(), - session.toSessionRepresentation(), - FAILED, - GENERAL_POOL, - false, - self, - ImmutableList.of(), - query, - immediateFailureQueryStats(), - Optional.empty(), - Optional.empty(), - Optional.empty(), - ImmutableMap.of(), - ImmutableSet.of(), - ImmutableMap.of(), - ImmutableMap.of(), - ImmutableSet.of(), - Optional.empty(), - false, - null, - Optional.empty(), - failureCause, - failureCause.getErrorCode(), - ImmutableList.of(), - ImmutableSet.of(), - Optional.empty(), - true, - resourceGroupId, - queryType); - - return queryInfo; - } - @JsonProperty public QueryId getQueryId() { diff --git a/presto-main/src/main/java/com/facebook/presto/execution/QueryManager.java b/presto-main/src/main/java/com/facebook/presto/execution/QueryManager.java index c498ac0638b6e..efcb622ab2bc3 100644 --- a/presto-main/src/main/java/com/facebook/presto/execution/QueryManager.java +++ b/presto-main/src/main/java/com/facebook/presto/execution/QueryManager.java @@ -13,10 +13,9 @@ */ package com.facebook.presto.execution; -import com.facebook.presto.execution.QueryExecution.QueryOutputInfo; +import com.facebook.presto.Session; import com.facebook.presto.execution.StateMachine.StateChangeListener; import com.facebook.presto.server.BasicQueryInfo; -import com.facebook.presto.server.SessionContext; import com.facebook.presto.spi.QueryId; import com.google.common.util.concurrent.ListenableFuture; @@ -33,7 +32,7 @@ public interface QueryManager * * @throws NoSuchElementException if query does not exist */ - void addOutputInfoListener(QueryId queryId, Consumer listener) + void addOutputInfoListener(QueryId queryId, Consumer listener) throws NoSuchElementException; /** @@ -66,6 +65,16 @@ BasicQueryInfo getQueryInfo(QueryId queryId) QueryInfo getFullQueryInfo(QueryId queryId) throws NoSuchElementException; + /** + * @throws NoSuchElementException if query does not exist + */ + Session getQuerySession(QueryId queryId); + + /** + * @throws NoSuchElementException if query does not exist + */ + boolean isQuerySlugValid(QueryId queryId, String slug); + /** * @throws NoSuchElementException if query does not exist */ @@ -78,13 +87,10 @@ QueryState getQueryState(QueryId queryId) */ void recordHeartbeat(QueryId queryId); - QueryId createQueryId(); - /** - * Creates a new query. This method may be called multiple times for the same query id. The - * the first call will be accepted, and the other calls will be ignored. + * Creates a new query using the specified query execution. */ - ListenableFuture createQuery(QueryId queryId, SessionContext sessionContext, String query); + void createQuery(QueryExecution execution); /** * Attempts to fail the query for the specified reason. If the query is already in a final @@ -104,5 +110,5 @@ QueryState getQueryState(QueryId queryId) */ void cancelStage(StageId stageId); - SqlQueryManagerStats getStats(); + QueryManagerStats getStats(); } diff --git a/presto-main/src/main/java/com/facebook/presto/execution/QueryManagerConfig.java b/presto-main/src/main/java/com/facebook/presto/execution/QueryManagerConfig.java index 29ff95dcde5a0..5596f2a69e23b 100644 --- a/presto-main/src/main/java/com/facebook/presto/execution/QueryManagerConfig.java +++ b/presto-main/src/main/java/com/facebook/presto/execution/QueryManagerConfig.java @@ -62,9 +62,6 @@ public class QueryManagerConfig private Duration queryMaxExecutionTime = new Duration(100, TimeUnit.DAYS); private Duration queryMaxCpuTime = new Duration(1_000_000_000, TimeUnit.DAYS); - private int initializationRequiredWorkers = 1; - private Duration initializationTimeout = new Duration(5, TimeUnit.MINUTES); - private int requiredWorkers = 1; private Duration requiredWorkersMaxWait = new Duration(5, TimeUnit.MINUTES); @@ -354,34 +351,6 @@ public QueryManagerConfig setQueryExecutionPolicy(String queryExecutionPolicy) return this; } - @Min(1) - public int getInitializationRequiredWorkers() - { - return initializationRequiredWorkers; - } - - @Config("query-manager.initialization-required-workers") - @ConfigDescription("Minimum number of workers that must be available before the cluster will accept queries") - public QueryManagerConfig setInitializationRequiredWorkers(int initializationRequiredWorkers) - { - this.initializationRequiredWorkers = initializationRequiredWorkers; - return this; - } - - @NotNull - public Duration getInitializationTimeout() - { - return initializationTimeout; - } - - @Config("query-manager.initialization-timeout") - @ConfigDescription("After this time, the cluster will accept queries even if the minimum required workers are not available") - public QueryManagerConfig setInitializationTimeout(Duration initializationTimeout) - { - this.initializationTimeout = initializationTimeout; - return this; - } - @Min(1) public int getRequiredWorkers() { diff --git a/presto-main/src/main/java/com/facebook/presto/execution/SqlQueryManagerStats.java b/presto-main/src/main/java/com/facebook/presto/execution/QueryManagerStats.java similarity index 74% rename from presto-main/src/main/java/com/facebook/presto/execution/SqlQueryManagerStats.java rename to presto-main/src/main/java/com/facebook/presto/execution/QueryManagerStats.java index 2f621a8bcfafe..39ffeb4b7382e 100644 --- a/presto-main/src/main/java/com/facebook/presto/execution/SqlQueryManagerStats.java +++ b/presto-main/src/main/java/com/facebook/presto/execution/QueryManagerStats.java @@ -13,20 +13,28 @@ */ package com.facebook.presto.execution; +import com.facebook.presto.dispatcher.DispatchQuery; +import com.facebook.presto.execution.StateMachine.StateChangeListener; +import com.facebook.presto.server.BasicQueryInfo; import io.airlift.stats.CounterStat; import io.airlift.stats.DistributionStat; import io.airlift.stats.TimeStat; import org.weakref.jmx.Managed; import org.weakref.jmx.Nested; +import javax.annotation.concurrent.GuardedBy; + +import java.util.Optional; import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Supplier; +import static com.facebook.presto.execution.QueryState.RUNNING; import static com.facebook.presto.spi.StandardErrorCode.ABANDONED_QUERY; import static com.facebook.presto.spi.StandardErrorCode.USER_CANCELED; import static java.util.concurrent.TimeUnit.MILLISECONDS; import static java.util.concurrent.TimeUnit.SECONDS; -public class SqlQueryManagerStats +public class QueryManagerStats { private final AtomicInteger queuedQueries = new AtomicInteger(); private final AtomicInteger runningQueries = new AtomicInteger(); @@ -48,25 +56,34 @@ public class SqlQueryManagerStats private final DistributionStat wallInputBytesRate = new DistributionStat(); private final DistributionStat cpuInputByteRate = new DistributionStat(); - public void queryQueued() + public void trackQueryStats(DispatchQuery managedQueryExecution) { submittedQueries.update(1); queuedQueries.incrementAndGet(); + managedQueryExecution.addStateChangeListener(new StatisticsListener(managedQueryExecution)); } - public void queryStarted() + public void trackQueryStats(QueryExecution managedQueryExecution) + { + submittedQueries.update(1); + queuedQueries.incrementAndGet(); + managedQueryExecution.addStateChangeListener(new StatisticsListener()); + managedQueryExecution.addFinalQueryInfoListener(finalQueryInfo -> queryFinished(new BasicQueryInfo(finalQueryInfo))); + } + + private void queryStarted() { startedQueries.update(1); runningQueries.incrementAndGet(); queuedQueries.decrementAndGet(); } - public void queryStopped() + private void queryStopped() { runningQueries.decrementAndGet(); } - public void queryFinished(QueryInfo info) + private void queryFinished(BasicQueryInfo info) { completedQueries.update(1); @@ -114,6 +131,52 @@ else if (info.getErrorCode().getCode() == USER_CANCELED.toErrorCode().getCode()) } } + private class StatisticsListener + implements StateChangeListener + { + private final Supplier> finalQueryInfoSupplier; + + @GuardedBy("this") + private boolean stopped; + @GuardedBy("this") + private boolean started; + + public StatisticsListener() + { + finalQueryInfoSupplier = Optional::empty; + } + + public StatisticsListener(DispatchQuery managedQueryExecution) + { + finalQueryInfoSupplier = () -> Optional.of(managedQueryExecution.getBasicQueryInfo()); + } + + @Override + public void stateChanged(QueryState newValue) + { + synchronized (this) { + if (stopped) { + return; + } + + if (newValue.isDone()) { + stopped = true; + if (started) { + queryStopped(); + } + finalQueryInfoSupplier.get() + .ifPresent(QueryManagerStats.this::queryFinished); + } + else if (newValue.ordinal() >= RUNNING.ordinal()) { + if (!started) { + started = true; + queryStarted(); + } + } + } + } + } + @Managed public long getRunningQueries() { diff --git a/presto-main/src/main/java/com/facebook/presto/execution/QueryState.java b/presto-main/src/main/java/com/facebook/presto/execution/QueryState.java index 59775e4022ec3..32fcd09264e05 100644 --- a/presto-main/src/main/java/com/facebook/presto/execution/QueryState.java +++ b/presto-main/src/main/java/com/facebook/presto/execution/QueryState.java @@ -28,6 +28,10 @@ public enum QueryState * Query is waiting for the required resources (beta). */ WAITING_FOR_RESOURCES(false), + /** + * Query is being dispatched to a coordinator. + */ + DISPATCHING(false), /** * Query is being planned. */ diff --git a/presto-main/src/main/java/com/facebook/presto/execution/QueryStateMachine.java b/presto-main/src/main/java/com/facebook/presto/execution/QueryStateMachine.java index ab9c501944886..511a2faadd6a1 100644 --- a/presto-main/src/main/java/com/facebook/presto/execution/QueryStateMachine.java +++ b/presto-main/src/main/java/com/facebook/presto/execution/QueryStateMachine.java @@ -68,6 +68,7 @@ import java.util.function.Consumer; import static com.facebook.presto.execution.BasicStageStats.EMPTY_STAGE_STATS; +import static com.facebook.presto.execution.QueryState.DISPATCHING; import static com.facebook.presto.execution.QueryState.FAILED; import static com.facebook.presto.execution.QueryState.FINISHED; import static com.facebook.presto.execution.QueryState.FINISHING; @@ -93,14 +94,14 @@ @ThreadSafe public class QueryStateMachine { - public static final Logger QUERY_STATE_LOG = Logger.get(QueryStateMachine.class); + private static final Logger QUERY_STATE_LOG = Logger.get(QueryStateMachine.class); private final QueryId queryId; private final String query; private final Session session; private final URI self; private final Optional queryType; - private final Optional resourceGroup; + private final ResourceGroupId resourceGroup; private final TransactionManager transactionManager; private final Metadata metadata; private final QueryOutputManager outputManager; @@ -150,7 +151,7 @@ private QueryStateMachine( String query, Session session, URI self, - Optional resourceGroup, + ResourceGroupId resourceGroup, Optional queryType, TransactionManager transactionManager, Executor executor, @@ -230,7 +231,7 @@ static QueryStateMachine beginWithTicker( query, session, self, - Optional.of(resourceGroup), + resourceGroup, queryType, transactionManager, executor, @@ -340,7 +341,7 @@ public BasicQueryInfo getBasicQueryInfo(Optional rootStage) return new BasicQueryInfo( queryId, session.toSessionRepresentation(), - resourceGroup, + Optional.of(resourceGroup), state, memoryPool.get().getId(), stageStats.isScheduled(), @@ -402,7 +403,7 @@ QueryInfo getQueryInfo(Optional rootStage) inputs.get(), output.get(), completeInfo, - resourceGroup, + Optional.of(resourceGroup), queryType); } @@ -520,6 +521,7 @@ private QueryStats getQueryStats(Optional rootStage) queryStateTimer.getElapsedTime(), queryStateTimer.getQueuedTime(), queryStateTimer.getResourceWaitingTime(), + queryStateTimer.getDispatchingTime(), queryStateTimer.getExecutionTime(), queryStateTimer.getAnalysisTime(), queryStateTimer.getPlanningTime(), @@ -713,6 +715,12 @@ public boolean transitionToWaitingForResources() return queryState.setIf(WAITING_FOR_RESOURCES, currentState -> currentState.ordinal() < WAITING_FOR_RESOURCES.ordinal()); } + public boolean transitionToDispatching() + { + queryStateTimer.beginDispatching(); + return queryState.setIf(DISPATCHING, currentState -> currentState.ordinal() < DISPATCHING.ordinal()); + } + public boolean transitionToPlanning() { queryStateTimer.beginPlanning(); @@ -996,6 +1004,7 @@ private static QueryStats pruneQueryStats(QueryStats queryStats) queryStats.getElapsedTime(), queryStats.getQueuedTime(), queryStats.getResourceWaitingTime(), + queryStats.getDispatchingTime(), queryStats.getExecutionTime(), queryStats.getAnalysisTime(), queryStats.getTotalPlanningTime(), diff --git a/presto-main/src/main/java/com/facebook/presto/execution/QueryStateTimer.java b/presto-main/src/main/java/com/facebook/presto/execution/QueryStateTimer.java index f343bca04c807..5390d0f6eafa4 100644 --- a/presto-main/src/main/java/com/facebook/presto/execution/QueryStateTimer.java +++ b/presto-main/src/main/java/com/facebook/presto/execution/QueryStateTimer.java @@ -33,12 +33,14 @@ class QueryStateTimer private final long createNanos; private final AtomicReference beginResourceWaitingNanos = new AtomicReference<>(); + private final AtomicReference beginDispatchingNanos = new AtomicReference<>(); private final AtomicReference beginPlanningNanos = new AtomicReference<>(); private final AtomicReference beginFinishingNanos = new AtomicReference<>(); private final AtomicReference endNanos = new AtomicReference<>(); private final AtomicReference queuedTime = new AtomicReference<>(); private final AtomicReference resourceWaitingTime = new AtomicReference<>(); + private final AtomicReference dispatchingTime = new AtomicReference<>(); private final AtomicReference executionTime = new AtomicReference<>(); private final AtomicReference planningTime = new AtomicReference<>(); private final AtomicReference finishingTime = new AtomicReference<>(); @@ -70,6 +72,18 @@ private void beginWaitingForResources(long now) beginResourceWaitingNanos.compareAndSet(null, now); } + public void beginDispatching() + { + beginDispatching(tickerNanos()); + } + + private void beginDispatching(long now) + { + beginWaitingForResources(now); + resourceWaitingTime.compareAndSet(null, nanosSince(beginResourceWaitingNanos, now)); + beginDispatchingNanos.compareAndSet(null, now); + } + public void beginPlanning() { beginPlanning(tickerNanos()); @@ -77,8 +91,8 @@ public void beginPlanning() private void beginPlanning(long now) { - beginWaitingForResources(now); - resourceWaitingTime.compareAndSet(null, nanosSince(beginResourceWaitingNanos, now)); + beginDispatching(now); + dispatchingTime.compareAndSet(null, nanosSince(beginDispatchingNanos, now)); beginPlanningNanos.compareAndSet(null, now); } @@ -184,6 +198,11 @@ public Duration getResourceWaitingTime() return getDuration(resourceWaitingTime, beginResourceWaitingNanos); } + public Duration getDispatchingTime() + { + return getDuration(dispatchingTime, beginDispatchingNanos); + } + public Duration getPlanningTime() { return getDuration(planningTime, beginPlanningNanos); @@ -234,7 +253,7 @@ private static Duration nanosSince(AtomicReference start, long end) private static Duration nanosSince(long start, long now) { - return succinctNanos(now - start); + return succinctNanos(Math.max(0, now - start)); } private Duration getDuration(AtomicReference finalDuration, AtomicReference start) diff --git a/presto-main/src/main/java/com/facebook/presto/execution/QueryStats.java b/presto-main/src/main/java/com/facebook/presto/execution/QueryStats.java index 5c716384ae43b..bbd8b3b4c71bf 100644 --- a/presto-main/src/main/java/com/facebook/presto/execution/QueryStats.java +++ b/presto-main/src/main/java/com/facebook/presto/execution/QueryStats.java @@ -31,11 +31,9 @@ import java.util.Set; import static com.google.common.base.Preconditions.checkArgument; -import static io.airlift.units.DataSize.Unit.BYTE; import static io.airlift.units.DataSize.succinctBytes; import static java.lang.Math.min; import static java.util.Objects.requireNonNull; -import static java.util.concurrent.TimeUnit.MILLISECONDS; public class QueryStats { @@ -47,8 +45,9 @@ public class QueryStats private final Duration elapsedTime; private final Duration queuedTime; - private final Duration executionTime; private final Duration resourceWaitingTime; + private final Duration dispatchingTime; + private final Duration executionTime; private final Duration analysisTime; private final Duration totalPlanningTime; private final Duration finishingTime; @@ -107,6 +106,7 @@ public QueryStats( @JsonProperty("elapsedTime") Duration elapsedTime, @JsonProperty("queuedTime") Duration queuedTime, @JsonProperty("resourceWaitingTime") Duration resourceWaitingTime, + @JsonProperty("dispatchingTime") Duration dispatchingTime, @JsonProperty("executionTime") Duration executionTime, @JsonProperty("analysisTime") Duration analysisTime, @JsonProperty("totalPlanningTime") Duration totalPlanningTime, @@ -164,6 +164,7 @@ public QueryStats( this.elapsedTime = requireNonNull(elapsedTime, "elapsedTime is null"); this.queuedTime = requireNonNull(queuedTime, "queuedTime is null"); this.resourceWaitingTime = requireNonNull(resourceWaitingTime, "resourceWaitingTime is null"); + this.dispatchingTime = requireNonNull(dispatchingTime, "dispatchingTime is null"); this.executionTime = requireNonNull(executionTime, "executionTime is null"); this.analysisTime = requireNonNull(analysisTime, "analysisTime is null"); this.totalPlanningTime = requireNonNull(totalPlanningTime, "totalPlanningTime is null"); @@ -224,56 +225,6 @@ public QueryStats( this.operatorSummaries = ImmutableList.copyOf(requireNonNull(operatorSummaries, "operatorSummaries is null")); } - public static QueryStats immediateFailureQueryStats() - { - DateTime now = DateTime.now(); - return new QueryStats( - now, - now, - now, - now, - new Duration(0, MILLISECONDS), - new Duration(0, MILLISECONDS), - new Duration(0, MILLISECONDS), - new Duration(0, MILLISECONDS), - new Duration(0, MILLISECONDS), - new Duration(0, MILLISECONDS), - new Duration(0, MILLISECONDS), - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - new DataSize(0, BYTE), - new DataSize(0, BYTE), - new DataSize(0, BYTE), - new DataSize(0, BYTE), - new DataSize(0, BYTE), - new DataSize(0, BYTE), - false, - new Duration(0, MILLISECONDS), - new Duration(0, MILLISECONDS), - new Duration(0, MILLISECONDS), - false, - ImmutableSet.of(), - new DataSize(0, BYTE), - 0, - new DataSize(0, BYTE), - 0, - new DataSize(0, BYTE), - 0, - 0, - new DataSize(0, BYTE), - new DataSize(0, BYTE), - new DataSize(0, BYTE), - ImmutableList.of(), - ImmutableList.of()); - } - @JsonProperty public DateTime getCreateTime() { @@ -311,6 +262,12 @@ public Duration getResourceWaitingTime() return resourceWaitingTime; } + @JsonProperty + public Duration getDispatchingTime() + { + return dispatchingTime; + } + @JsonProperty public Duration getQueuedTime() { diff --git a/presto-main/src/main/java/com/facebook/presto/execution/QueryTracker.java b/presto-main/src/main/java/com/facebook/presto/execution/QueryTracker.java index 9ef467d8f78ad..227059fbe3c14 100644 --- a/presto-main/src/main/java/com/facebook/presto/execution/QueryTracker.java +++ b/presto-main/src/main/java/com/facebook/presto/execution/QueryTracker.java @@ -144,7 +144,7 @@ public T getQuery(QueryId queryId) throws NoSuchElementException { return tryGetQuery(queryId) - .orElseThrow(NoSuchElementException::new); + .orElseThrow(() -> new NoSuchElementException(queryId.toString())); } public Optional tryGetQuery(QueryId queryId) diff --git a/presto-main/src/main/java/com/facebook/presto/execution/SqlQueryExecution.java b/presto-main/src/main/java/com/facebook/presto/execution/SqlQueryExecution.java index aca1d1ada241d..3c502f8c17d82 100644 --- a/presto-main/src/main/java/com/facebook/presto/execution/SqlQueryExecution.java +++ b/presto-main/src/main/java/com/facebook/presto/execution/SqlQueryExecution.java @@ -34,11 +34,9 @@ import com.facebook.presto.operator.ForScheduler; import com.facebook.presto.security.AccessControl; import com.facebook.presto.server.BasicQueryInfo; -import com.facebook.presto.spi.ErrorCode; import com.facebook.presto.spi.PrestoException; import com.facebook.presto.spi.QueryId; import com.facebook.presto.spi.resourceGroups.QueryType; -import com.facebook.presto.spi.resourceGroups.ResourceGroupId; import com.facebook.presto.split.CloseableSplitSourceProvider; import com.facebook.presto.split.SplitManager; import com.facebook.presto.sql.analyzer.Analysis; @@ -59,7 +57,6 @@ import com.facebook.presto.sql.planner.optimizations.PlanOptimizer; import com.facebook.presto.sql.planner.plan.OutputNode; import com.facebook.presto.sql.tree.Explain; -import com.facebook.presto.transaction.TransactionManager; import com.google.common.collect.ImmutableSet; import com.google.common.util.concurrent.ListenableFuture; import io.airlift.concurrent.SetThreadName; @@ -70,7 +67,6 @@ import javax.annotation.concurrent.ThreadSafe; import javax.inject.Inject; -import java.net.URI; import java.util.List; import java.util.Map; import java.util.Optional; @@ -86,8 +82,6 @@ import static com.facebook.presto.spi.StandardErrorCode.NOT_SUPPORTED; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Throwables.throwIfInstanceOf; -import static io.airlift.concurrent.MoreFutures.addExceptionCallback; -import static io.airlift.concurrent.MoreFutures.addSuccessCallback; import static io.airlift.units.DataSize.Unit.BYTE; import static io.airlift.units.DataSize.succinctBytes; import static java.util.Objects.requireNonNull; @@ -100,8 +94,7 @@ public class SqlQueryExecution private static final OutputBufferId OUTPUT_BUFFER_ID = new OutputBufferId(0); private final QueryStateMachine stateMachine; - - private final ClusterSizeMonitor clusterSizeMonitor; + private final String slug; private final Metadata metadata; private final SqlParser sqlParser; private final SplitManager splitManager; @@ -126,14 +119,9 @@ public class SqlQueryExecution private final CostCalculator costCalculator; private SqlQueryExecution( - String query, - Session session, - URI self, - ResourceGroupId resourceGroup, - Optional queryType, PreparedQuery preparedQuery, - ClusterSizeMonitor clusterSizeMonitor, - TransactionManager transactionManager, + QueryStateMachine stateMachine, + String slug, Metadata metadata, AccessControl accessControl, SqlParser sqlParser, @@ -156,8 +144,8 @@ private SqlQueryExecution( CostCalculator costCalculator, WarningCollector warningCollector) { - try (SetThreadName ignored = new SetThreadName("Query-%s", session.getQueryId())) { - this.clusterSizeMonitor = requireNonNull(clusterSizeMonitor, "clusterSizeMonitor is null"); + try (SetThreadName ignored = new SetThreadName("Query-%s", stateMachine.getQueryId())) { + this.slug = requireNonNull(slug, "slug is null"); this.metadata = requireNonNull(metadata, "metadata is null"); this.sqlParser = requireNonNull(sqlParser, "sqlParser is null"); this.splitManager = requireNonNull(splitManager, "splitManager is null"); @@ -178,21 +166,7 @@ private SqlQueryExecution( checkArgument(scheduleSplitBatchSize > 0, "scheduleSplitBatchSize must be greater than 0"); this.scheduleSplitBatchSize = scheduleSplitBatchSize; - requireNonNull(query, "query is null"); - requireNonNull(session, "session is null"); - requireNonNull(self, "self is null"); - this.stateMachine = QueryStateMachine.begin( - query, - session, - self, - resourceGroup, - queryType, - false, - transactionManager, - accessControl, - queryExecutor, - metadata, - warningCollector); + this.stateMachine = requireNonNull(stateMachine, "stateMachine is null"); // analyze query requireNonNull(preparedQuery, "preparedQuery is null"); @@ -233,6 +207,12 @@ private SqlQueryExecution( } } + @Override + public String getSlug() + { + return slug; + } + @Override public VersionedMemoryPoolId getMemoryPool() { @@ -327,20 +307,6 @@ public BasicQueryInfo getBasicQueryInfo() @Override public void start() - { - if (stateMachine.transitionToWaitingForResources()) { - waitForMinimumWorkers(); - } - } - - private void waitForMinimumWorkers() - { - ListenableFuture minimumWorkerFuture = clusterSizeMonitor.waitForMinimumWorkers(); - addSuccessCallback(minimumWorkerFuture, () -> queryExecutor.submit(this::startExecution)); - addExceptionCallback(minimumWorkerFuture, throwable -> queryExecutor.submit(() -> stateMachine.transitionToFailed(throwable))); - } - - private void startExecution() { try (SetThreadName ignored = new SetThreadName("Query-%s", stateMachine.getQueryId())) { try { @@ -392,12 +358,6 @@ public Session getSession() return stateMachine.getSession(); } - @Override - public Optional getErrorCode() - { - return stateMachine.getFailureInfo().map(ExecutionFailureInfo::getErrorCode); - } - @Override public void addFinalQueryInfoListener(StateChangeListener stateChangeListener) { @@ -663,7 +623,6 @@ public static class SqlQueryExecutionFactory private final List planOptimizers; private final PlanFragmenter planFragmenter; private final RemoteTaskFactory remoteTaskFactory; - private final TransactionManager transactionManager; private final QueryExplainer queryExplainer; private final LocationFactory locationFactory; private final ExecutorService queryExecutor; @@ -671,7 +630,6 @@ public static class SqlQueryExecutionFactory private final FailureDetector failureDetector; private final NodeTaskMap nodeTaskMap; private final Map executionPolicies; - private final ClusterSizeMonitor clusterSizeMonitor; private final StatsCalculator statsCalculator; private final CostCalculator costCalculator; @@ -687,7 +645,6 @@ public static class SqlQueryExecutionFactory PlanOptimizers planOptimizers, PlanFragmenter planFragmenter, RemoteTaskFactory remoteTaskFactory, - TransactionManager transactionManager, @ForQueryExecution ExecutorService queryExecutor, @ForScheduler ScheduledExecutorService schedulerExecutor, FailureDetector failureDetector, @@ -695,7 +652,6 @@ public static class SqlQueryExecutionFactory QueryExplainer queryExplainer, Map executionPolicies, SplitSchedulerStats schedulerStats, - ClusterSizeMonitor clusterSizeMonitor, StatsCalculator statsCalculator, CostCalculator costCalculator) { @@ -712,14 +668,12 @@ public static class SqlQueryExecutionFactory requireNonNull(planOptimizers, "planOptimizers is null"); this.planFragmenter = requireNonNull(planFragmenter, "planFragmenter is null"); this.remoteTaskFactory = requireNonNull(remoteTaskFactory, "remoteTaskFactory is null"); - this.transactionManager = requireNonNull(transactionManager, "transactionManager is null"); this.queryExecutor = requireNonNull(queryExecutor, "queryExecutor is null"); this.schedulerExecutor = requireNonNull(schedulerExecutor, "schedulerExecutor is null"); this.failureDetector = requireNonNull(failureDetector, "failureDetector is null"); this.nodeTaskMap = requireNonNull(nodeTaskMap, "nodeTaskMap is null"); this.queryExplainer = requireNonNull(queryExplainer, "queryExplainer is null"); this.executionPolicies = requireNonNull(executionPolicies, "schedulerPolicies is null"); - this.clusterSizeMonitor = requireNonNull(clusterSizeMonitor, "clusterSizeMonitor is null"); this.planOptimizers = planOptimizers.get(); this.statsCalculator = requireNonNull(statsCalculator, "statsCalculator is null"); this.costCalculator = requireNonNull(costCalculator, "costCalculator is null"); @@ -727,26 +681,20 @@ public static class SqlQueryExecutionFactory @Override public QueryExecution createQueryExecution( - String query, - Session session, PreparedQuery preparedQuery, - ResourceGroupId resourceGroup, + QueryStateMachine stateMachine, + String slug, WarningCollector warningCollector, Optional queryType) { - String executionPolicyName = SystemSessionProperties.getExecutionPolicy(session); + String executionPolicyName = SystemSessionProperties.getExecutionPolicy(stateMachine.getSession()); ExecutionPolicy executionPolicy = executionPolicies.get(executionPolicyName); checkArgument(executionPolicy != null, "No execution policy %s", executionPolicy); SqlQueryExecution execution = new SqlQueryExecution( - query, - session, - locationFactory.createQueryLocation(session.getQueryId()), - resourceGroup, - queryType, preparedQuery, - clusterSizeMonitor, - transactionManager, + stateMachine, + slug, metadata, accessControl, sqlParser, diff --git a/presto-main/src/main/java/com/facebook/presto/execution/SqlQueryManager.java b/presto-main/src/main/java/com/facebook/presto/execution/SqlQueryManager.java index adf7664b307f8..61e2b007bb629 100644 --- a/presto-main/src/main/java/com/facebook/presto/execution/SqlQueryManager.java +++ b/presto-main/src/main/java/com/facebook/presto/execution/SqlQueryManager.java @@ -16,34 +16,16 @@ import com.facebook.presto.ExceededCpuLimitException; import com.facebook.presto.Session; import com.facebook.presto.event.QueryMonitor; -import com.facebook.presto.execution.QueryExecution.QueryExecutionFactory; import com.facebook.presto.execution.QueryExecution.QueryOutputInfo; -import com.facebook.presto.execution.QueryPreparer.PreparedQuery; import com.facebook.presto.execution.StateMachine.StateChangeListener; -import com.facebook.presto.execution.resourceGroups.ResourceGroupManager; -import com.facebook.presto.execution.scheduler.NodeSchedulerConfig; -import com.facebook.presto.execution.warnings.WarningCollector; import com.facebook.presto.execution.warnings.WarningCollectorFactory; import com.facebook.presto.memory.ClusterMemoryManager; -import com.facebook.presto.metadata.SessionPropertyManager; -import com.facebook.presto.security.AccessControl; import com.facebook.presto.server.BasicQueryInfo; -import com.facebook.presto.server.SessionContext; -import com.facebook.presto.server.SessionPropertyDefaults; -import com.facebook.presto.server.SessionSupplier; import com.facebook.presto.spi.PrestoException; import com.facebook.presto.spi.QueryId; -import com.facebook.presto.spi.resourceGroups.QueryType; -import com.facebook.presto.spi.resourceGroups.SelectionContext; -import com.facebook.presto.spi.resourceGroups.SelectionCriteria; -import com.facebook.presto.sql.SqlEnvironmentConfig; -import com.facebook.presto.sql.SqlPath; import com.facebook.presto.sql.planner.Plan; -import com.facebook.presto.sql.tree.Statement; -import com.facebook.presto.transaction.TransactionManager; import com.facebook.presto.version.EmbedVersion; import com.google.common.collect.Ordering; -import com.google.common.util.concurrent.AbstractFuture; import com.google.common.util.concurrent.ListenableFuture; import io.airlift.concurrent.ThreadPoolExecutorMBean; import io.airlift.log.Logger; @@ -58,27 +40,18 @@ import javax.inject.Inject; import java.util.List; -import java.util.Map; import java.util.NoSuchElementException; import java.util.Objects; -import java.util.Optional; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.Consumer; import static com.facebook.presto.SystemSessionProperties.getQueryMaxCpuTime; import static com.facebook.presto.execution.QueryState.RUNNING; -import static com.facebook.presto.execution.QueryStateMachine.QUERY_STATE_LOG; -import static com.facebook.presto.spi.StandardErrorCode.NOT_SUPPORTED; -import static com.facebook.presto.spi.StandardErrorCode.QUERY_TEXT_TOO_LARGE; -import static com.facebook.presto.util.Failures.toFailure; -import static com.facebook.presto.util.StatementUtils.getQueryType; -import static com.facebook.presto.util.StatementUtils.isTransactionControlStatement; -import static com.google.common.base.Preconditions.checkArgument; +import static com.facebook.presto.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR; import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.util.concurrent.Futures.immediateFailedFuture; import static io.airlift.concurrent.Threads.threadsNamed; @@ -92,96 +65,35 @@ public class SqlQueryManager { private static final Logger log = Logger.get(SqlQueryManager.class); - private final QueryPreparer queryPreparer; - - private final EmbedVersion embedVersion; - private final ExecutorService queryExecutor; - private final ThreadPoolExecutorMBean queryExecutorMBean; - private final ResourceGroupManager resourceGroupManager; private final ClusterMemoryManager memoryManager; + private final QueryMonitor queryMonitor; + private final EmbedVersion embedVersion; + private final QueryTracker queryTracker; - private final Optional path; - private final int maxQueryLength; private final Duration maxQueryCpuTime; - private final QueryTracker queryTracker; + private final ExecutorService queryExecutor; + private final ThreadPoolExecutorMBean queryExecutorMBean; private final ScheduledExecutorService queryManagementExecutor; private final ThreadPoolExecutorMBean queryManagementExecutorMBean; - private final QueryMonitor queryMonitor; - private final LocationFactory locationFactory; - - private final TransactionManager transactionManager; - private final AccessControl accessControl; - - private final QueryIdGenerator queryIdGenerator; - - private final SessionSupplier sessionSupplier; - private final SessionPropertyDefaults sessionPropertyDefaults; - - private final ClusterSizeMonitor clusterSizeMonitor; - - private final Map, QueryExecutionFactory> executionFactories; - - private final SqlQueryManagerStats stats = new SqlQueryManagerStats(); - - private final WarningCollectorFactory warningCollectorFactory; + private final QueryManagerStats stats = new QueryManagerStats(); @Inject - public SqlQueryManager( - QueryPreparer queryPreparer, - EmbedVersion embedVersion, - NodeSchedulerConfig nodeSchedulerConfig, - QueryManagerConfig queryManagerConfig, - SqlEnvironmentConfig sqlEnvironmentConfig, - QueryMonitor queryMonitor, - ResourceGroupManager resourceGroupManager, - ClusterMemoryManager memoryManager, - LocationFactory locationFactory, - TransactionManager transactionManager, - AccessControl accessControl, - QueryIdGenerator queryIdGenerator, - SessionSupplier sessionSupplier, - SessionPropertyDefaults sessionPropertyDefaults, - ClusterSizeMonitor clusterSizeMonitor, - Map, QueryExecutionFactory> executionFactories, - WarningCollectorFactory warningCollectorFactory) + public SqlQueryManager(ClusterMemoryManager memoryManager, QueryMonitor queryMonitor, EmbedVersion embedVersion, QueryManagerConfig queryManagerConfig, WarningCollectorFactory warningCollectorFactory) { - this.queryPreparer = requireNonNull(queryPreparer, "queryPreparer is null"); - - this.embedVersion = requireNonNull(embedVersion, "embedVersion is null"); - this.executionFactories = requireNonNull(executionFactories, "executionFactories is null"); - - this.queryExecutor = newCachedThreadPool(threadsNamed("query-scheduler-%s")); - this.queryExecutorMBean = new ThreadPoolExecutorMBean((ThreadPoolExecutor) queryExecutor); - - requireNonNull(nodeSchedulerConfig, "nodeSchedulerConfig is null"); - requireNonNull(queryManagerConfig, "queryManagerConfig is null"); - this.resourceGroupManager = requireNonNull(resourceGroupManager, "resourceGroupManager is null"); this.memoryManager = requireNonNull(memoryManager, "memoryManager is null"); - this.queryMonitor = requireNonNull(queryMonitor, "queryMonitor is null"); - this.locationFactory = requireNonNull(locationFactory, "locationFactory is null"); - - this.transactionManager = requireNonNull(transactionManager, "transactionManager is null"); - this.accessControl = requireNonNull(accessControl, "accessControl is null"); - - this.queryIdGenerator = requireNonNull(queryIdGenerator, "queryIdGenerator is null"); - - this.sessionSupplier = requireNonNull(sessionSupplier, "sessionSupplier is null"); - this.sessionPropertyDefaults = requireNonNull(sessionPropertyDefaults, "sessionPropertyDefaults is null"); - - this.clusterSizeMonitor = requireNonNull(clusterSizeMonitor, "clusterSizeMonitor is null"); + this.embedVersion = requireNonNull(embedVersion, "embedVersion is null"); - this.path = sqlEnvironmentConfig.getPath(); - this.maxQueryLength = queryManagerConfig.getMaxQueryLength(); this.maxQueryCpuTime = queryManagerConfig.getQueryMaxCpuTime(); - this.warningCollectorFactory = requireNonNull(warningCollectorFactory, "warningCollectorFactory is null"); + this.queryExecutor = newCachedThreadPool(threadsNamed("query-scheduler-%s")); + this.queryExecutorMBean = new ThreadPoolExecutorMBean((ThreadPoolExecutor) queryExecutor); - queryManagementExecutor = Executors.newScheduledThreadPool(queryManagerConfig.getQueryManagerExecutorPoolSize(), threadsNamed("query-management-%s")); - queryManagementExecutorMBean = new ThreadPoolExecutorMBean((ThreadPoolExecutor) queryManagementExecutor); + this.queryManagementExecutor = Executors.newScheduledThreadPool(queryManagerConfig.getQueryManagerExecutorPoolSize(), threadsNamed("query-management-%s")); + this.queryManagementExecutorMBean = new ThreadPoolExecutorMBean((ThreadPoolExecutor) queryManagementExecutor); this.queryTracker = new QueryTracker<>(queryManagerConfig, queryManagementExecutor); } @@ -263,10 +175,24 @@ public BasicQueryInfo getQueryInfo(QueryId queryId) @Override public QueryInfo getFullQueryInfo(QueryId queryId) + throws NoSuchElementException { return queryTracker.getQuery(queryId).getQueryInfo(); } + @Override + public Session getQuerySession(QueryId queryId) + throws NoSuchElementException + { + return queryTracker.getQuery(queryId).getSession(); + } + + @Override + public boolean isQuerySlugValid(QueryId queryId, String slug) + { + return queryTracker.getQuery(queryId).getSlug().equals(slug); + } + public Plan getQueryPlan(QueryId queryId) { return queryTracker.getQuery(queryId).getQueryPlan(); @@ -291,158 +217,27 @@ public void recordHeartbeat(QueryId queryId) } @Override - public QueryId createQueryId() + public void createQuery(QueryExecution queryExecution) { - return queryIdGenerator.createNextQueryId(); - } + requireNonNull(queryExecution, "queryExecution is null"); - @Override - public ListenableFuture createQuery(QueryId queryId, SessionContext sessionContext, String query) - { - QueryCreationFuture queryCreationFuture = new QueryCreationFuture(); - queryExecutor.submit(embedVersion.embedVersion(() -> { - try { - createQueryInternal(queryId, sessionContext, query, resourceGroupManager); - queryCreationFuture.set(null); - } - catch (Throwable e) { - queryCreationFuture.setException(e); - } - })); - return queryCreationFuture; - } - - private void createQueryInternal(QueryId queryId, SessionContext sessionContext, String query, ResourceGroupManager resourceGroupManager) - { - requireNonNull(queryId, "queryId is null"); - requireNonNull(sessionContext, "sessionFactory is null"); - requireNonNull(query, "query is null"); - checkArgument(!query.isEmpty(), "query must not be empty string"); - checkArgument(!queryTracker.tryGetQuery(queryId).isPresent(), "query %s already exists", queryId); - - Session session = null; - SelectionContext selectionContext = null; - QueryExecution queryExecution; - PreparedQuery preparedQuery; - Optional queryType = Optional.empty(); - try { - clusterSizeMonitor.verifyInitialMinimumWorkersRequirement(); - - if (query.length() > maxQueryLength) { - int queryLength = query.length(); - query = query.substring(0, maxQueryLength); - throw new PrestoException(QUERY_TEXT_TOO_LARGE, format("Query text length (%s) exceeds the maximum length (%s)", queryLength, maxQueryLength)); - } - - // decode session - session = sessionSupplier.createSession(queryId, sessionContext); - - WarningCollector warningCollector = warningCollectorFactory.create(); - - // prepare query - preparedQuery = queryPreparer.prepareQuery(session, query, warningCollector); - - // select resource group - queryType = getQueryType(preparedQuery.getStatement().getClass()); - selectionContext = resourceGroupManager.selectGroup(new SelectionCriteria( - sessionContext.getIdentity().getPrincipal().isPresent(), - sessionContext.getIdentity().getUser(), - Optional.ofNullable(sessionContext.getSource()), - sessionContext.getClientTags(), - sessionContext.getResourceEstimates(), - queryType.map(Enum::name))); - - // apply system defaults for query - session = sessionPropertyDefaults.newSessionWithDefaultProperties(session, queryType.map(Enum::name), selectionContext.getResourceGroupId()); - - // mark existing transaction as active - transactionManager.activateTransaction(session, isTransactionControlStatement(preparedQuery.getStatement()), accessControl); - - // create query execution - QueryExecutionFactory queryExecutionFactory = executionFactories.get(preparedQuery.getStatement().getClass()); - if (queryExecutionFactory == null) { - throw new PrestoException(NOT_SUPPORTED, "Unsupported statement type: " + preparedQuery.getStatement().getClass().getSimpleName()); - } - queryExecution = queryExecutionFactory.createQueryExecution( - query, - session, - preparedQuery, - selectionContext.getResourceGroupId(), - warningCollector, - queryType); - } - catch (RuntimeException e) { - // This is intentionally not a method, since after the state change listener is registered - // it's not safe to do any of this, and we had bugs before where people reused this code in a method - - // if session creation failed, create a minimal session object - if (session == null) { - session = Session.builder(new SessionPropertyManager()) - .setQueryId(queryId) - .setIdentity(sessionContext.getIdentity()) - .setPath(new SqlPath(Optional.empty())) - .build(); - } - QUERY_STATE_LOG.debug(e, "Query %s failed", session.getQueryId()); - - // query failure fails the transaction - session.getTransactionId().ifPresent(transactionManager::fail); - - QueryExecution execution = new FailedQueryExecution( - session, - query, - locationFactory.createQueryLocation(queryId), - Optional.ofNullable(selectionContext).map(SelectionContext::getResourceGroupId), - queryType, - queryExecutor, - e); - - try { - queryTracker.addQuery(execution); - - BasicQueryInfo queryInfo = execution.getBasicQueryInfo(); - queryMonitor.queryCreatedEvent(queryInfo); - queryMonitor.queryImmediateFailureEvent(queryInfo, toFailure(e)); - stats.queryQueued(); - stats.queryStarted(); - stats.queryStopped(); - stats.queryFinished(execution.getQueryInfo()); - } - finally { - // execution MUST be added to the expiration queue or there will be a leak - queryTracker.expireQuery(queryId); - } - - return; + if (!queryTracker.addQuery(queryExecution)) { + throw new PrestoException(GENERIC_INTERNAL_ERROR, format("Query %s already registered", queryExecution.getQueryId())); } - queryMonitor.queryCreatedEvent(queryExecution.getBasicQueryInfo()); - queryExecution.addFinalQueryInfoListener(finalQueryInfo -> { try { - stats.queryFinished(finalQueryInfo); queryMonitor.queryCompletedEvent(finalQueryInfo); } finally { // execution MUST be added to the expiration queue or there will be a leak - queryTracker.expireQuery(queryId); + queryTracker.expireQuery(queryExecution.getQueryId()); } }); - addStatsListeners(queryExecution); + stats.trackQueryStats(queryExecution); - if (!queryTracker.addQuery(queryExecution)) { - // query already created, so just exit - return; - } - - // start the query in the background - try { - resourceGroupManager.submit(preparedQuery.getStatement(), queryExecution, selectionContext, queryExecutor); - } - catch (Throwable e) { - failQuery(queryId, e); - } + embedVersion.embedVersion(queryExecution::start).run(); } @Override @@ -477,7 +272,7 @@ public void cancelStage(StageId stageId) @Override @Managed @Flatten - public SqlQueryManagerStats getStats() + public QueryManagerStats getStats() { return stats; } @@ -521,53 +316,4 @@ private void enforceCpuLimits() } } } - - private void addStatsListeners(QueryExecution queryExecution) - { - Object lock = new Object(); - - // QUEUED is the initial state, the counter can be incremented immediately - stats.queryQueued(); - - AtomicBoolean started = new AtomicBoolean(); - queryExecution.addStateChangeListener(newValue -> { - synchronized (lock) { - if (newValue == RUNNING && !started.getAndSet(true)) { - stats.queryStarted(); - } - } - }); - - AtomicBoolean stopped = new AtomicBoolean(); - queryExecution.addStateChangeListener(newValue -> { - synchronized (lock) { - if (newValue.isDone() && !stopped.getAndSet(true) && started.get()) { - stats.queryStopped(); - } - } - }); - } - - private static class QueryCreationFuture - extends AbstractFuture - { - @Override - protected boolean set(QueryInfo value) - { - return super.set(value); - } - - @Override - protected boolean setException(Throwable throwable) - { - return super.setException(throwable); - } - - @Override - public boolean cancel(boolean mayInterruptIfRunning) - { - // query submission can not be canceled - return false; - } - } } diff --git a/presto-main/src/main/java/com/facebook/presto/execution/resourceGroups/InternalResourceGroup.java b/presto-main/src/main/java/com/facebook/presto/execution/resourceGroups/InternalResourceGroup.java index 79e9f7033e4ff..71c1f463d8695 100644 --- a/presto-main/src/main/java/com/facebook/presto/execution/resourceGroups/InternalResourceGroup.java +++ b/presto-main/src/main/java/com/facebook/presto/execution/resourceGroups/InternalResourceGroup.java @@ -662,7 +662,7 @@ private void startInBackground(ManagedQueryExecution query) group = group.parent.get(); } updateEligibility(); - executor.execute(query::start); + executor.execute(query::startWaitingForResources); } } diff --git a/presto-main/src/main/java/com/facebook/presto/server/BasicQueryInfo.java b/presto-main/src/main/java/com/facebook/presto/server/BasicQueryInfo.java index 0927bb175e117..6ecb19bbb7512 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/BasicQueryInfo.java +++ b/presto-main/src/main/java/com/facebook/presto/server/BasicQueryInfo.java @@ -13,6 +13,7 @@ */ package com.facebook.presto.server; +import com.facebook.presto.Session; import com.facebook.presto.SessionRepresentation; import com.facebook.presto.execution.QueryInfo; import com.facebook.presto.execution.QueryState; @@ -25,6 +26,7 @@ import com.facebook.presto.spi.resourceGroups.ResourceGroupId; import com.fasterxml.jackson.annotation.JsonCreator; import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.common.collect.ImmutableList; import javax.annotation.Nullable; import javax.annotation.concurrent.Immutable; @@ -33,6 +35,9 @@ import java.util.List; import java.util.Optional; +import static com.facebook.presto.execution.QueryState.FAILED; +import static com.facebook.presto.memory.LocalMemoryManager.GENERAL_POOL; +import static com.facebook.presto.server.BasicQueryStats.immediateFailureQueryStats; import static com.google.common.base.MoreObjects.toStringHelper; import static java.util.Objects.requireNonNull; @@ -105,6 +110,24 @@ public BasicQueryInfo(QueryInfo queryInfo) queryInfo.getWarnings()); } + public static BasicQueryInfo immediateFailureQueryInfo(Session session, String query, URI self, Optional resourceGroupId, ErrorCode errorCode) + { + return new BasicQueryInfo( + session.getQueryId(), + session.toSessionRepresentation(), + resourceGroupId, + FAILED, + GENERAL_POOL, + false, + self, + query, + immediateFailureQueryStats(), + errorCode == null ? null : errorCode.getType(), + errorCode, + Optional.empty(), + ImmutableList.of()); + } + @JsonProperty public QueryId getQueryId() { diff --git a/presto-main/src/main/java/com/facebook/presto/server/BasicQueryStats.java b/presto-main/src/main/java/com/facebook/presto/server/BasicQueryStats.java index a395f7f8bf093..dbd5e8f326712 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/BasicQueryStats.java +++ b/presto-main/src/main/java/com/facebook/presto/server/BasicQueryStats.java @@ -28,7 +28,9 @@ import java.util.Set; import static com.google.common.base.Preconditions.checkArgument; +import static io.airlift.units.DataSize.Unit.BYTE; import static java.util.Objects.requireNonNull; +import static java.util.concurrent.TimeUnit.MILLISECONDS; /** * Lightweight version of QueryStats. Parts of the web UI depend on the fields @@ -143,6 +145,32 @@ public BasicQueryStats(QueryStats queryStats) queryStats.getProgressPercentage()); } + public static BasicQueryStats immediateFailureQueryStats() + { + DateTime now = DateTime.now(); + return new BasicQueryStats( + now, + now, + new Duration(0, MILLISECONDS), + new Duration(0, MILLISECONDS), + new Duration(0, MILLISECONDS), + 0, + 0, + 0, + 0, + new DataSize(0, BYTE), + 0, + 0, + new DataSize(0, BYTE), + new DataSize(0, BYTE), + new DataSize(0, BYTE), + new Duration(0, MILLISECONDS), + new Duration(0, MILLISECONDS), + false, + ImmutableSet.of(), + OptionalDouble.empty()); + } + @JsonProperty public DateTime getCreateTime() { diff --git a/presto-main/src/main/java/com/facebook/presto/server/CoordinatorModule.java b/presto-main/src/main/java/com/facebook/presto/server/CoordinatorModule.java index 96379cdea20ff..f0fdf0d59d663 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/CoordinatorModule.java +++ b/presto-main/src/main/java/com/facebook/presto/server/CoordinatorModule.java @@ -21,6 +21,12 @@ import com.facebook.presto.cost.CostComparator; import com.facebook.presto.cost.StatsCalculatorModule; import com.facebook.presto.cost.TaskCountEstimator; +import com.facebook.presto.dispatcher.DispatchExecutor; +import com.facebook.presto.dispatcher.DispatchManager; +import com.facebook.presto.dispatcher.DispatchQueryFactory; +import com.facebook.presto.dispatcher.FailedDispatchQueryFactory; +import com.facebook.presto.dispatcher.LocalDispatchQueryFactory; +import com.facebook.presto.dispatcher.QueuedStatementResource; import com.facebook.presto.event.QueryMonitor; import com.facebook.presto.event.QueryMonitorConfig; import com.facebook.presto.execution.AddColumnTask; @@ -84,7 +90,7 @@ import com.facebook.presto.memory.TotalReservationOnBlockedNodesLowMemoryKiller; import com.facebook.presto.metadata.CatalogManager; import com.facebook.presto.operator.ForScheduler; -import com.facebook.presto.server.protocol.StatementResource; +import com.facebook.presto.server.protocol.ExecutingStatementResource; import com.facebook.presto.server.remotetask.HttpRemoteTaskFactory; import com.facebook.presto.server.remotetask.RemoteTaskStats; import com.facebook.presto.spi.memory.ClusterMemoryPoolManager; @@ -187,8 +193,8 @@ protected void setup(Binder binder) jsonCodecBinder(binder).bindJsonCodec(TaskInfo.class); jsonCodecBinder(binder).bindJsonCodec(QueryResults.class); jsonCodecBinder(binder).bindJsonCodec(SelectedRole.class); - jaxrsBinder(binder).bind(StatementResource.class); - newExporter(binder).export(StatementResource.class).withGeneratedName(); + jaxrsBinder(binder).bind(QueuedStatementResource.class); + jaxrsBinder(binder).bind(ExecutingStatementResource.class); binder.bind(StatementHttpExecutionMBean.class).in(Scopes.SINGLETON); newExporter(binder).export(StatementHttpExecutionMBean.class).withGeneratedName(); @@ -212,13 +218,21 @@ protected void setup(Binder binder) jaxrsBinder(binder).bind(ResourceGroupStateInfoResource.class); binder.bind(QueryIdGenerator.class).in(Scopes.SINGLETON); binder.bind(QueryManager.class).to(SqlQueryManager.class).in(Scopes.SINGLETON); + newExporter(binder).export(QueryManager.class).withGeneratedName(); binder.bind(QueryPreparer.class).in(Scopes.SINGLETON); binder.bind(SessionSupplier.class).to(QuerySessionSupplier.class).in(Scopes.SINGLETON); binder.bind(InternalResourceGroupManager.class).in(Scopes.SINGLETON); newExporter(binder).export(InternalResourceGroupManager.class).withGeneratedName(); binder.bind(ResourceGroupManager.class).to(InternalResourceGroupManager.class); binder.bind(LegacyResourceGroupConfigurationManager.class).in(Scopes.SINGLETON); - newExporter(binder).export(QueryManager.class).withGeneratedName(); + + // dispatcher + binder.bind(DispatchManager.class).in(Scopes.SINGLETON); + binder.bind(FailedDispatchQueryFactory.class).in(Scopes.SINGLETON); + binder.bind(DispatchExecutor.class).in(Scopes.SINGLETON); + + // local dispatcher + binder.bind(DispatchQueryFactory.class).to(LocalDispatchQueryFactory.class); // cluster memory manager binder.bind(ClusterMemoryManager.class).in(Scopes.SINGLETON); diff --git a/presto-main/src/main/java/com/facebook/presto/server/QueryResource.java b/presto-main/src/main/java/com/facebook/presto/server/QueryResource.java index adc0ff78f2969..b3f56789657bc 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/QueryResource.java +++ b/presto-main/src/main/java/com/facebook/presto/server/QueryResource.java @@ -13,6 +13,7 @@ */ package com.facebook.presto.server; +import com.facebook.presto.dispatcher.DispatchManager; import com.facebook.presto.execution.QueryInfo; import com.facebook.presto.execution.QueryManager; import com.facebook.presto.execution.QueryState; @@ -45,11 +46,14 @@ @Path("/v1/query") public class QueryResource { + // TODO There should be a combined interface for this + private final DispatchManager dispatchManager; private final QueryManager queryManager; @Inject - public QueryResource(QueryManager queryManager) + public QueryResource(DispatchManager dispatchManager, QueryManager queryManager) { + this.dispatchManager = requireNonNull(dispatchManager, "dispatchManager is null"); this.queryManager = requireNonNull(queryManager, "queryManager is null"); } @@ -58,7 +62,7 @@ public List getAllQueryInfo(@QueryParam("state") String stateFil { QueryState expectedState = stateFilter == null ? null : QueryState.valueOf(stateFilter.toUpperCase(Locale.ENGLISH)); ImmutableList.Builder builder = new ImmutableList.Builder<>(); - for (BasicQueryInfo queryInfo : queryManager.getQueries()) { + for (BasicQueryInfo queryInfo : dispatchManager.getQueries()) { if (stateFilter == null || queryInfo.getState() == expectedState) { builder.add(queryInfo); } diff --git a/presto-main/src/main/java/com/facebook/presto/server/QueryStateInfoResource.java b/presto-main/src/main/java/com/facebook/presto/server/QueryStateInfoResource.java index 1f38c8d932b62..38e21b8fe0af8 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/QueryStateInfoResource.java +++ b/presto-main/src/main/java/com/facebook/presto/server/QueryStateInfoResource.java @@ -13,7 +13,7 @@ */ package com.facebook.presto.server; -import com.facebook.presto.execution.QueryManager; +import com.facebook.presto.dispatcher.DispatchManager; import com.facebook.presto.execution.resourceGroups.ResourceGroupManager; import com.facebook.presto.spi.QueryId; import com.facebook.presto.spi.resourceGroups.ResourceGroupId; @@ -43,15 +43,15 @@ @Path("/v1/queryState") public class QueryStateInfoResource { - private final QueryManager queryManager; + private final DispatchManager dispatchManager; private final ResourceGroupManager resourceGroupManager; @Inject public QueryStateInfoResource( - QueryManager queryManager, + DispatchManager dispatchManager, ResourceGroupManager resourceGroupManager) { - this.queryManager = requireNonNull(queryManager, "queryManager is null"); + this.dispatchManager = requireNonNull(dispatchManager, "dispatchManager is null"); this.resourceGroupManager = requireNonNull(resourceGroupManager, "resourceGroupManager is null"); } @@ -59,7 +59,7 @@ public QueryStateInfoResource( @Produces(MediaType.APPLICATION_JSON) public List getQueryStateInfos(@QueryParam("user") String user) { - List queryInfos = queryManager.getQueries(); + List queryInfos = dispatchManager.getQueries(); if (!isNullOrEmpty(user)) { queryInfos = queryInfos.stream() @@ -92,7 +92,7 @@ public QueryStateInfo getQueryStateInfo(@PathParam("queryId") String queryId) throws WebApplicationException { try { - return getQueryStateInfo(queryManager.getQueryInfo(new QueryId(queryId))); + return getQueryStateInfo(dispatchManager.getQueryInfo(new QueryId(queryId))); } catch (NoSuchElementException e) { throw new WebApplicationException(NOT_FOUND); diff --git a/presto-main/src/main/java/com/facebook/presto/server/protocol/StatementResource.java b/presto-main/src/main/java/com/facebook/presto/server/protocol/ExecutingStatementResource.java similarity index 69% rename from presto-main/src/main/java/com/facebook/presto/server/protocol/StatementResource.java rename to presto-main/src/main/java/com/facebook/presto/server/protocol/ExecutingStatementResource.java index 3f10205c2e6db..238b0f008387f 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/protocol/StatementResource.java +++ b/presto-main/src/main/java/com/facebook/presto/server/protocol/ExecutingStatementResource.java @@ -13,34 +13,26 @@ */ package com.facebook.presto.server.protocol; +import com.facebook.presto.Session; import com.facebook.presto.client.QueryResults; import com.facebook.presto.execution.QueryManager; import com.facebook.presto.memory.context.SimpleLocalMemoryContext; -import com.facebook.presto.metadata.SessionPropertyManager; import com.facebook.presto.operator.ExchangeClient; import com.facebook.presto.operator.ExchangeClientSupplier; import com.facebook.presto.server.ForStatementResource; -import com.facebook.presto.server.HttpRequestSessionContext; -import com.facebook.presto.server.SessionContext; import com.facebook.presto.spi.QueryId; import com.facebook.presto.spi.block.BlockEncodingSerde; import com.google.common.collect.Ordering; import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.ListenableFuture; import io.airlift.concurrent.BoundedExecutor; -import io.airlift.stats.CounterStat; import io.airlift.units.DataSize; import io.airlift.units.Duration; -import org.weakref.jmx.Managed; -import org.weakref.jmx.Nested; -import javax.annotation.PreDestroy; import javax.inject.Inject; -import javax.servlet.http.HttpServletRequest; import javax.ws.rs.DELETE; import javax.ws.rs.GET; import javax.ws.rs.HeaderParam; -import javax.ws.rs.POST; import javax.ws.rs.Path; import javax.ws.rs.PathParam; import javax.ws.rs.Produces; @@ -58,7 +50,7 @@ import java.io.UnsupportedEncodingException; import java.net.URLEncoder; import java.util.Map.Entry; -import java.util.OptionalLong; +import java.util.NoSuchElementException; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; import java.util.concurrent.ScheduledExecutorService; @@ -77,16 +69,15 @@ import static com.google.common.base.Strings.isNullOrEmpty; import static com.google.common.net.HttpHeaders.X_FORWARDED_PROTO; import static com.google.common.util.concurrent.MoreExecutors.directExecutor; -import static io.airlift.concurrent.Threads.threadsNamed; import static io.airlift.http.server.AsyncResponseHandler.bindAsyncResponse; import static io.airlift.units.DataSize.Unit.MEGABYTE; import static java.util.Objects.requireNonNull; -import static java.util.concurrent.Executors.newSingleThreadScheduledExecutor; -import static java.util.concurrent.TimeUnit.MILLISECONDS; import static java.util.concurrent.TimeUnit.SECONDS; +import static javax.ws.rs.core.MediaType.TEXT_PLAIN_TYPE; +import static javax.ws.rs.core.Response.Status.NOT_FOUND; -@Path("/v1/statement") -public class StatementResource +@Path("/") +public class ExecutingStatementResource { private static final Duration MAX_WAIT_TIME = new Duration(1, SECONDS); private static final Ordering> WAIT_ORDERING = Ordering.natural().nullsLast(); @@ -95,86 +86,34 @@ public class StatementResource private static final DataSize MAX_TARGET_RESULT_SIZE = new DataSize(128, MEGABYTE); private final QueryManager queryManager; - private final SessionPropertyManager sessionPropertyManager; private final ExchangeClientSupplier exchangeClientSupplier; private final BlockEncodingSerde blockEncodingSerde; private final BoundedExecutor responseExecutor; private final ScheduledExecutorService timeoutExecutor; private final ConcurrentMap queries = new ConcurrentHashMap<>(); - private final ScheduledExecutorService queryPurger = newSingleThreadScheduledExecutor(threadsNamed("query-purger")); - - private final CounterStat createQueryRequests = new CounterStat(); @Inject - public StatementResource( + public ExecutingStatementResource( QueryManager queryManager, - SessionPropertyManager sessionPropertyManager, ExchangeClientSupplier exchangeClientSupplier, BlockEncodingSerde blockEncodingSerde, @ForStatementResource BoundedExecutor responseExecutor, @ForStatementResource ScheduledExecutorService timeoutExecutor) { this.queryManager = requireNonNull(queryManager, "queryManager is null"); - this.sessionPropertyManager = requireNonNull(sessionPropertyManager, "sessionPropertyManager is null"); this.exchangeClientSupplier = requireNonNull(exchangeClientSupplier, "exchangeClientSupplier is null"); this.blockEncodingSerde = requireNonNull(blockEncodingSerde, "blockEncodingSerde is null"); this.responseExecutor = requireNonNull(responseExecutor, "responseExecutor is null"); this.timeoutExecutor = requireNonNull(timeoutExecutor, "timeoutExecutor is null"); - - queryPurger.scheduleWithFixedDelay(new PurgeQueriesRunnable(queries, queryManager), 200, 200, MILLISECONDS); - } - - @PreDestroy - public void stop() - { - queryPurger.shutdownNow(); - } - - @POST - @Produces(MediaType.APPLICATION_JSON) - public Response createQuery( - String statement, - @HeaderParam(X_FORWARDED_PROTO) String proto, - @Context HttpServletRequest servletRequest, - @Context UriInfo uriInfo) - { - createQueryRequests.update(1); - - if (isNullOrEmpty(statement)) { - throw new WebApplicationException(Response - .status(Status.BAD_REQUEST) - .type(MediaType.TEXT_PLAIN) - .entity("SQL statement is empty") - .build()); - } - if (isNullOrEmpty(proto)) { - proto = uriInfo.getRequestUri().getScheme(); - } - - SessionContext sessionContext = new HttpRequestSessionContext(servletRequest); - - ExchangeClient exchangeClient = exchangeClientSupplier.get(new SimpleLocalMemoryContext(newSimpleAggregatedMemoryContext(), StatementResource.class.getSimpleName())); - Query query = Query.create( - sessionContext, - statement, - queryManager, - sessionPropertyManager, - exchangeClient, - responseExecutor, - timeoutExecutor, - blockEncodingSerde); - queries.put(query.getQueryId(), query); - - QueryResults queryResults = query.getNextResult(OptionalLong.empty(), uriInfo, proto, DEFAULT_TARGET_RESULT_SIZE); - return toResponse(query, queryResults); } @GET - @Path("{queryId}/{token}") + @Path("/v1/statement/executing/{queryId}/{slug}/{token}") @Produces(MediaType.APPLICATION_JSON) public void getQueryResults( @PathParam("queryId") QueryId queryId, + @PathParam("slug") String slug, @PathParam("token") long token, @QueryParam("maxWait") Duration maxWait, @QueryParam("targetResultSize") DataSize targetResultSize, @@ -182,21 +121,53 @@ public void getQueryResults( @Context UriInfo uriInfo, @Suspended AsyncResponse asyncResponse) { - Query query = queries.get(queryId); - if (query == null) { - asyncResponse.resume(Response.status(Status.NOT_FOUND).build()); - return; - } + Query query = getQuery(queryId, slug); if (isNullOrEmpty(proto)) { proto = uriInfo.getRequestUri().getScheme(); } - asyncQueryResults(query, OptionalLong.of(token), maxWait, targetResultSize, uriInfo, proto, asyncResponse); + asyncQueryResults(query, token, maxWait, targetResultSize, uriInfo, proto, asyncResponse); + } + + protected Query getQuery(QueryId queryId, String slug) + { + Query query = queries.get(queryId); + if (query != null) { + if (!query.isSlugValid(slug)) { + throw badRequest(NOT_FOUND, "Query not found"); + } + return query; + } + + // this is the first time the query has been accessed on this coordinator + Session session; + try { + if (!queryManager.isQuerySlugValid(queryId, slug)) { + throw badRequest(NOT_FOUND, "Query not found"); + } + session = queryManager.getQuerySession(queryId); + } + catch (NoSuchElementException e) { + throw badRequest(NOT_FOUND, "Query not found"); + } + + query = queries.computeIfAbsent(queryId, id -> { + ExchangeClient exchangeClient = exchangeClientSupplier.get(new SimpleLocalMemoryContext(newSimpleAggregatedMemoryContext(), ExecutingStatementResource.class.getSimpleName())); + return Query.create( + session, + slug, + queryManager, + exchangeClient, + responseExecutor, + timeoutExecutor, + blockEncodingSerde); + }); + return query; } private void asyncQueryResults( Query query, - OptionalLong token, + long token, Duration maxWait, DataSize targetResultSize, UriInfo uriInfo, @@ -226,16 +197,16 @@ private static Response toResponse(Query query, QueryResults queryResults) query.getSetPath().ifPresent(path -> response.header(PRESTO_SET_PATH, path)); // add set session properties - query.getSetSessionProperties().entrySet() - .forEach(entry -> response.header(PRESTO_SET_SESSION, entry.getKey() + '=' + entry.getValue())); + query.getSetSessionProperties() + .forEach((key, value) -> response.header(PRESTO_SET_SESSION, key + '=' + urlEncode(value))); // add clear session properties query.getResetSessionProperties() .forEach(name -> response.header(PRESTO_CLEAR_SESSION, name)); // add set roles - query.getSetRoles().entrySet() - .forEach(entry -> response.header(PRESTO_SET_ROLE, entry.getKey() + '=' + urlEncode(entry.getValue().toString()))); + query.getSetRoles() + .forEach((key, value) -> response.header(PRESTO_SET_ROLE, key + '=' + urlEncode(value.toString()))); // add added prepare statements for (Entry entry : query.getAddedPreparedStatements().entrySet()) { @@ -262,17 +233,42 @@ private static Response toResponse(Query query, QueryResults queryResults) } @DELETE - @Path("{queryId}/{token}") + @Path("/v1/statement/executing/{queryId}/{slug}/{token}") @Produces(MediaType.APPLICATION_JSON) - public Response cancelQuery(@PathParam("queryId") QueryId queryId, + public Response cancelQuery( + @PathParam("queryId") QueryId queryId, + @PathParam("slug") String slug, @PathParam("token") long token) { Query query = queries.get(queryId); - if (query == null) { - return Response.status(Status.NOT_FOUND).build(); + if (query != null) { + if (!query.isSlugValid(slug)) { + throw badRequest(NOT_FOUND, "Query not found"); + } + query.cancel(); + return Response.noContent().build(); } - query.cancel(); - return Response.noContent().build(); + + // cancel the query execution directly instead of creating the statement client + try { + if (!queryManager.isQuerySlugValid(queryId, slug)) { + throw badRequest(NOT_FOUND, "Query not found"); + } + queryManager.cancelQuery(queryId); + return Response.noContent().build(); + } + catch (NoSuchElementException e) { + throw badRequest(NOT_FOUND, "Query not found"); + } + } + + private static WebApplicationException badRequest(Status status, String message) + { + throw new WebApplicationException( + Response.status(status) + .type(TEXT_PLAIN_TYPE) + .entity(message) + .build()); } private static String urlEncode(String value) @@ -284,11 +280,4 @@ private static String urlEncode(String value) throw new AssertionError(e); } } - - @Managed - @Nested - public CounterStat getCreateQueryRequests() - { - return createQueryRequests; - } } diff --git a/presto-main/src/main/java/com/facebook/presto/server/protocol/PurgeQueriesRunnable.java b/presto-main/src/main/java/com/facebook/presto/server/protocol/PurgeQueriesRunnable.java deleted file mode 100644 index 07af5a3e2a1f4..0000000000000 --- a/presto-main/src/main/java/com/facebook/presto/server/protocol/PurgeQueriesRunnable.java +++ /dev/null @@ -1,71 +0,0 @@ -/* - * 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 com.facebook.presto.server.protocol; - -import com.facebook.presto.execution.QueryManager; -import com.facebook.presto.execution.QueryState; -import com.facebook.presto.spi.QueryId; -import com.google.common.collect.ImmutableSet; -import io.airlift.log.Logger; - -import java.util.NoSuchElementException; -import java.util.concurrent.ConcurrentMap; - -class PurgeQueriesRunnable - implements Runnable -{ - private static final Logger log = Logger.get(PurgeQueriesRunnable.class); - - private final ConcurrentMap queries; - private final QueryManager queryManager; - - public PurgeQueriesRunnable(ConcurrentMap queries, QueryManager queryManager) - { - this.queries = queries; - this.queryManager = queryManager; - } - - @Override - public void run() - { - try { - // Queries are added to the query manager before being recorded in queryIds set. - // Therefore, we take a snapshot if queryIds before getting the live queries - // from the query manager. Then we remove only the queries in the snapshot and - // not live queries set. If we did this in the other order, a query could be - // registered between fetching the live queries and inspecting the queryIds set. - for (QueryId queryId : ImmutableSet.copyOf(queries.keySet())) { - Query query = queries.get(queryId); - if (!query.isSubmissionFinished()) { - continue; - } - try { - // free up resources if the query completed - if (queryManager.getQueryState(queryId) == QueryState.FAILED) { - query.dispose(); - } - } - catch (NoSuchElementException e) { - // make sure all resource are released - query.dispose(); - // forget about this query if the query manager is no longer tracking it - queries.remove(queryId); - } - } - } - catch (Throwable e) { - log.warn(e, "Error removing old queries"); - } - } -} diff --git a/presto-main/src/main/java/com/facebook/presto/server/protocol/Query.java b/presto-main/src/main/java/com/facebook/presto/server/protocol/Query.java index 76fe95c0eddbe..3023a751dd5cf 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/protocol/Query.java +++ b/presto-main/src/main/java/com/facebook/presto/server/protocol/Query.java @@ -30,9 +30,7 @@ import com.facebook.presto.execution.buffer.PagesSerde; import com.facebook.presto.execution.buffer.PagesSerdeFactory; import com.facebook.presto.execution.buffer.SerializedPage; -import com.facebook.presto.metadata.SessionPropertyManager; import com.facebook.presto.operator.ExchangeClient; -import com.facebook.presto.server.SessionContext; import com.facebook.presto.spi.ErrorCode; import com.facebook.presto.spi.Page; import com.facebook.presto.spi.QueryId; @@ -47,8 +45,6 @@ import com.google.common.collect.ImmutableSet; import com.google.common.collect.Iterables; import com.google.common.collect.Lists; -import com.google.common.util.concurrent.AbstractFuture; -import com.google.common.util.concurrent.FutureCallback; import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.ListenableFuture; import io.airlift.log.Logger; @@ -72,7 +68,6 @@ import java.util.Set; import java.util.concurrent.Executor; import java.util.concurrent.ScheduledExecutorService; -import java.util.concurrent.atomic.AtomicLong; import static com.facebook.presto.SystemSessionProperties.isExchangeCompressionEnabled; import static com.facebook.presto.execution.QueryState.FAILED; @@ -80,9 +75,9 @@ import static com.facebook.presto.util.Failures.toFailure; import static com.google.common.base.MoreObjects.firstNonNull; import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Verify.verify; import static com.google.common.util.concurrent.Futures.immediateFuture; import static com.google.common.util.concurrent.MoreExecutors.directExecutor; -import static io.airlift.concurrent.MoreFutures.addSuccessCallback; import static io.airlift.concurrent.MoreFutures.addTimeout; import static java.lang.String.format; import static java.util.Objects.requireNonNull; @@ -94,6 +89,8 @@ class Query private final QueryManager queryManager; private final QueryId queryId; + private final Session session; + private final String slug; @GuardedBy("this") private final ExchangeClient exchangeClient; @@ -101,23 +98,16 @@ class Query private final Executor resultsProcessorExecutor; private final ScheduledExecutorService timeoutExecutor; - @GuardedBy("this") - private PagesSerde serde; - - private final AtomicLong resultId = new AtomicLong(); - - private final QuerySubmissionFuture submissionFuture; - private final SessionPropertyManager sessionPropertyManager; - private final BlockEncodingSerde blockEncodingSerde; + private final PagesSerde serde; @GuardedBy("this") - private Session session; + private OptionalLong nextToken = OptionalLong.of(0); @GuardedBy("this") private QueryResults lastResult; @GuardedBy("this") - private String lastResultPath; + private long lastToken = -1; @GuardedBy("this") private List columns; @@ -159,76 +149,60 @@ class Query private Long updateCount; public static Query create( - SessionContext sessionContext, - String query, + Session session, + String slug, QueryManager queryManager, - SessionPropertyManager sessionPropertyManager, ExchangeClient exchangeClient, Executor dataProcessorExecutor, ScheduledExecutorService timeoutExecutor, BlockEncodingSerde blockEncodingSerde) { - Query result = new Query(sessionContext, query, queryManager, sessionPropertyManager, exchangeClient, dataProcessorExecutor, timeoutExecutor, blockEncodingSerde); + Query result = new Query(session, slug, queryManager, exchangeClient, dataProcessorExecutor, timeoutExecutor, blockEncodingSerde); - // register listeners after submission finishes - addSuccessCallback(result.submissionFuture, () -> { - result.queryManager.addOutputInfoListener(result.getQueryId(), result::setQueryOutputInfo); + result.queryManager.addOutputInfoListener(result.getQueryId(), result::setQueryOutputInfo); - result.queryManager.addStateChangeListener(result.getQueryId(), state -> { - if (state.isDone()) { - QueryInfo queryInfo = queryManager.getFullQueryInfo(result.getQueryId()); - result.closeExchangeClientIfNecessary(queryInfo); - } - }); + result.queryManager.addStateChangeListener(result.getQueryId(), state -> { + if (state.isDone()) { + QueryInfo queryInfo = queryManager.getFullQueryInfo(result.getQueryId()); + result.closeExchangeClientIfNecessary(queryInfo); + } }); return result; } private Query( - SessionContext sessionContext, - String query, + Session session, + String slug, QueryManager queryManager, - SessionPropertyManager sessionPropertyManager, ExchangeClient exchangeClient, Executor resultsProcessorExecutor, ScheduledExecutorService timeoutExecutor, BlockEncodingSerde blockEncodingSerde) { - requireNonNull(sessionContext, "sessionContext is null"); - requireNonNull(query, "query is null"); + requireNonNull(session, "session is null"); + requireNonNull(slug, "slug is null"); requireNonNull(queryManager, "queryManager is null"); - requireNonNull(sessionPropertyManager, "sessionPropertyManager is null"); requireNonNull(exchangeClient, "exchangeClient is null"); requireNonNull(resultsProcessorExecutor, "resultsProcessorExecutor is null"); requireNonNull(timeoutExecutor, "timeoutExecutor is null"); requireNonNull(blockEncodingSerde, "serde is null"); this.queryManager = queryManager; - this.sessionPropertyManager = sessionPropertyManager; - queryId = queryManager.createQueryId(); - submissionFuture = new QuerySubmissionFuture(queryId, query, sessionContext, queryManager); + this.queryId = session.getQueryId(); + this.session = session; + this.slug = slug; this.exchangeClient = exchangeClient; this.resultsProcessorExecutor = resultsProcessorExecutor; this.timeoutExecutor = timeoutExecutor; - this.blockEncodingSerde = blockEncodingSerde; - } - public boolean isSubmissionFinished() - { - return submissionFuture.isDone(); + serde = new PagesSerdeFactory(blockEncodingSerde, isExchangeCompressionEnabled(session)).createPagesSerde(); } public void cancel() { - // if submission is not finished, send cancel after it is finished - if (submissionFuture.isDone()) { - submissionFuture.addListener(() -> queryManager.cancelQuery(queryId), resultsProcessorExecutor); - } - else { - queryManager.cancelQuery(queryId); - } + queryManager.cancelQuery(queryId); dispose(); } @@ -242,6 +216,11 @@ public QueryId getQueryId() return queryId; } + public boolean isSlugValid(String slug) + { + return this.slug.equals(slug); + } + public synchronized Optional getSetCatalog() { return setCatalog; @@ -292,14 +271,12 @@ public synchronized boolean isClearTransactionId() return clearTransactionId; } - public synchronized ListenableFuture waitForResults(OptionalLong token, UriInfo uriInfo, String scheme, Duration wait, DataSize targetResultSize) + public synchronized ListenableFuture waitForResults(long token, UriInfo uriInfo, String scheme, Duration wait, DataSize targetResultSize) { // before waiting, check if this request has already been processed and cached - if (token.isPresent()) { - Optional cachedResult = getCachedResult(token.getAsLong(), uriInfo); - if (cachedResult.isPresent()) { - return immediateFuture(cachedResult.get()); - } + Optional cachedResult = getCachedResult(token); + if (cachedResult.isPresent()) { + return immediateFuture(cachedResult.get()); } // wait for a results data or query to finish, up to the wait timeout @@ -315,14 +292,6 @@ public synchronized ListenableFuture waitForResults(OptionalLong t private synchronized ListenableFuture getFutureStateChange() { - // ensure the query has been submitted - submissionFuture.submitQuery(); - - // if query query submission has not finished, wait for it to finish - if (!submissionFuture.isDone()) { - return submissionFuture; - } - // if the exchange client is open, wait for data if (!exchangeClient.isClosed()) { return exchangeClient.isBlocked(); @@ -338,24 +307,32 @@ private synchronized ListenableFuture getFutureStateChange() } } - private synchronized Optional getCachedResult(long token, UriInfo uriInfo) + private synchronized Optional getCachedResult(long token) { + // is this the first request? + if (lastResult == null) { + return Optional.empty(); + } + // is the a repeated request for the last results? - String requestedPath = uriInfo.getAbsolutePath().getPath(); - if (requestedPath.equals(lastResultPath)) { - if (submissionFuture.isDone()) { - // tell query manager we are still interested in the query - queryManager.recordHeartbeat(queryId); - } + if (token == lastToken) { + // tell query manager we are still interested in the query + queryManager.recordHeartbeat(queryId); return Optional.of(lastResult); } - if (token < resultId.get()) { + // if this is a result before the lastResult, the data is gone + if (token < lastToken) { throw new WebApplicationException(Response.Status.GONE); } + // if this is a request for a result after the end of the stream, return not found + if (!nextToken.isPresent()) { + throw new WebApplicationException(Response.Status.NOT_FOUND); + } + // if this is not a request for the next results, return not found - if (lastResult.getNextUri() == null || !requestedPath.equals(lastResult.getNextUri().getPath())) { + if (token != nextToken.getAsLong()) { // unknown token throw new WebApplicationException(Response.Status.NOT_FOUND); } @@ -363,50 +340,22 @@ private synchronized Optional getCachedResult(long token, UriInfo return Optional.empty(); } - public synchronized QueryResults getNextResult(OptionalLong token, UriInfo uriInfo, String scheme, DataSize targetResultSize) + private synchronized QueryResults getNextResult(long token, UriInfo uriInfo, String scheme, DataSize targetResultSize) { // check if the result for the token have already been created - if (token.isPresent()) { - Optional cachedResult = getCachedResult(token.getAsLong(), uriInfo); - if (cachedResult.isPresent()) { - return cachedResult.get(); - } + Optional cachedResult = getCachedResult(token); + if (cachedResult.isPresent()) { + return cachedResult.get(); } + verify(nextToken.isPresent(), "Can not generate next result when next token is not present"); + verify(token == nextToken.getAsLong(), "Expected token to equal next token"); URI queryHtmlUri = uriInfo.getRequestUriBuilder() .scheme(scheme) .replacePath("ui/query.html") .replaceQuery(queryId.toString()) .build(); - // if query query submission has not finished, return simple empty result - if (!submissionFuture.isDone()) { - QueryResults queryResults = new QueryResults( - queryId.toString(), - queryHtmlUri, - null, - createNextResultsUri(scheme, uriInfo), - null, - null, - StatementStats.builder() - .setState(QueryState.QUEUED.toString()) - .setQueued(true) - .setScheduled(false) - .build(), - null, - ImmutableList.of(), - null, - null); - - cacheLastResults(queryResults); - return queryResults; - } - - if (session == null) { - session = queryManager.getFullQueryInfo(queryId).getSession().toSession(sessionPropertyManager); - serde = new PagesSerdeFactory(blockEncodingSerde, isExchangeCompressionEnabled(session)).createPagesSerde(); - } - // Remove as many pages as possible from the exchange until just greater than DESIRED_RESULT_BYTES // NOTE: it is critical that query results are created for the pages removed from the exchange // client while holding the lock because the query may transition to the finished state when the @@ -464,14 +413,21 @@ public synchronized QueryResults getNextResult(OptionalLong token, UriInfo uriIn data = ImmutableSet.of(ImmutableList.of(true)); } + // advance next token // only return a next if // (1) the query is not done AND the query state is not FAILED // OR // (2)there is more data to send (due to buffering) + if ((!queryInfo.isFinalQueryInfo() && queryInfo.getState() != FAILED) || !exchangeClient.isClosed()) { + nextToken = OptionalLong.of(token + 1); + } + else { + nextToken = OptionalLong.empty(); + } + URI nextResultsUri = null; - if (!queryInfo.isFinalQueryInfo() && !queryInfo.getState().equals(QueryState.FAILED) - || !exchangeClient.isClosed()) { - nextResultsUri = createNextResultsUri(scheme, uriInfo); + if (nextToken.isPresent()) { + nextResultsUri = createNextResultsUri(scheme, uriInfo, nextToken.getAsLong()); } // update catalog, schema, and path @@ -508,20 +464,11 @@ public synchronized QueryResults getNextResult(OptionalLong token, UriInfo uriIn queryInfo.getUpdateType(), updateCount); - cacheLastResults(queryResults); - return queryResults; - } - - private synchronized void cacheLastResults(QueryResults queryResults) - { - // cache the last results - if (lastResult != null && lastResult.getNextUri() != null) { - lastResultPath = lastResult.getNextUri().getPath(); - } - else { - lastResultPath = null; - } + // cache the new result + lastToken = token; lastResult = queryResults; + + return queryResults; } private synchronized void closeExchangeClientIfNecessary(QueryInfo queryInfo) @@ -565,13 +512,14 @@ private ListenableFuture queryDoneFuture(QueryState currentState) return Futures.transformAsync(queryManager.getStateChange(queryId, currentState), this::queryDoneFuture, directExecutor()); } - private synchronized URI createNextResultsUri(String scheme, UriInfo uriInfo) + private synchronized URI createNextResultsUri(String scheme, UriInfo uriInfo, long nextToken) { return uriInfo.getBaseUriBuilder() .scheme(scheme) - .replacePath("/v1/statement") + .replacePath("/v1/statement/executing") .path(queryId.toString()) - .path(String.valueOf(resultId.incrementAndGet())) + .path(slug) + .path(String.valueOf(nextToken)) .replaceQuery("") .build(); } @@ -607,7 +555,6 @@ private static StageStats toStageStats(StageInfo stageInfo) if (stageInfo == null) { return null; } - com.facebook.presto.execution.StageStats stageStats = stageInfo.getStageStats(); ImmutableList.Builder subStages = ImmutableList.builder(); @@ -716,54 +663,4 @@ private static QueryError toQueryError(QueryInfo queryInfo) failure.getErrorLocation(), failure); } - - private static class QuerySubmissionFuture - extends AbstractFuture - { - private final QueryId queryId; - private final String query; - private final SessionContext sessionContext; - private final QueryManager queryManager; - - @GuardedBy("this") - private ListenableFuture querySubmissionFuture; - - public QuerySubmissionFuture(QueryId queryId, String query, SessionContext sessionContext, QueryManager queryManager) - { - this.queryId = requireNonNull(queryId, "queryId is null"); - this.query = requireNonNull(query, "query is null"); - this.sessionContext = requireNonNull(sessionContext, "sessionContext is null"); - this.queryManager = requireNonNull(queryManager, "queryManager is null"); - } - - private synchronized void submitQuery() - { - if (querySubmissionFuture != null) { - return; - } - - querySubmissionFuture = queryManager.createQuery(queryId, sessionContext, this.query); - Futures.addCallback(querySubmissionFuture, new FutureCallback() - { - @Override - public void onSuccess(Object result) - { - set(null); - } - - @Override - public void onFailure(Throwable t) - { - setException(t); - } - }, directExecutor()); - } - - @Override - public boolean cancel(boolean mayInterruptIfRunning) - { - // query submission can not be canceled - return false; - } - } } diff --git a/presto-main/src/main/java/com/facebook/presto/server/testing/TestingPrestoServer.java b/presto-main/src/main/java/com/facebook/presto/server/testing/TestingPrestoServer.java index 66d68f7f5e42f..713254f95a9ec 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/testing/TestingPrestoServer.java +++ b/presto-main/src/main/java/com/facebook/presto/server/testing/TestingPrestoServer.java @@ -16,6 +16,7 @@ import com.facebook.presto.connector.ConnectorId; import com.facebook.presto.connector.ConnectorManager; import com.facebook.presto.cost.StatsCalculator; +import com.facebook.presto.dispatcher.DispatchManager; import com.facebook.presto.eventlistener.EventListenerManager; import com.facebook.presto.execution.QueryInfo; import com.facebook.presto.execution.QueryManager; @@ -124,7 +125,7 @@ public class TestingPrestoServer private final StatsCalculator statsCalculator; private final TestingAccessControlManager accessControl; private final ProcedureTester procedureTester; - private final Optional resourceGroupManager; + private final Optional> resourceGroupManager; private final SplitManager splitManager; private final PageSourceManager pageSourceManager; private final NodePartitioningManager nodePartitioningManager; @@ -133,6 +134,7 @@ public class TestingPrestoServer private final InternalNodeManager nodeManager; private final ServiceSelectorManager serviceSelectorManager; private final Announcer announcer; + private final DispatchManager dispatchManager; private final SqlQueryManager queryManager; private final TaskManager taskManager; private final GracefulShutdownHandler gracefulShutdownHandler; @@ -282,13 +284,6 @@ public TestingPrestoServer( lifeCycleManager = injector.getInstance(LifeCycleManager.class); - if (coordinator) { - queryManager = (SqlQueryManager) injector.getInstance(QueryManager.class); - } - else { - queryManager = null; - } - pluginManager = injector.getInstance(PluginManager.class); connectorManager = injector.getInstance(ConnectorManager.class); @@ -302,12 +297,16 @@ public TestingPrestoServer( splitManager = injector.getInstance(SplitManager.class); pageSourceManager = injector.getInstance(PageSourceManager.class); if (coordinator) { + dispatchManager = injector.getInstance(DispatchManager.class); + queryManager = (SqlQueryManager) injector.getInstance(QueryManager.class); resourceGroupManager = Optional.of(injector.getInstance(InternalResourceGroupManager.class)); nodePartitioningManager = injector.getInstance(NodePartitioningManager.class); clusterMemoryManager = injector.getInstance(ClusterMemoryManager.class); statsCalculator = injector.getInstance(StatsCalculator.class); } else { + dispatchManager = null; + queryManager = null; resourceGroupManager = Optional.empty(); nodePartitioningManager = null; clusterMemoryManager = null; @@ -351,6 +350,11 @@ public void installPlugin(Plugin plugin) pluginManager.installPlugin(plugin); } + public DispatchManager getDispatchManager() + { + return dispatchManager; + } + public QueryManager getQueryManager() { return queryManager; @@ -445,7 +449,7 @@ public PageSourceManager getPageSourceManager() return pageSourceManager; } - public Optional getResourceGroupManager() + public Optional> getResourceGroupManager() { return resourceGroupManager; } diff --git a/presto-main/src/test/java/com/facebook/presto/execution/MockManagedQueryExecution.java b/presto-main/src/test/java/com/facebook/presto/execution/MockManagedQueryExecution.java new file mode 100644 index 0000000000000..38f59fa8f7728 --- /dev/null +++ b/presto-main/src/test/java/com/facebook/presto/execution/MockManagedQueryExecution.java @@ -0,0 +1,193 @@ +/* + * 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 com.facebook.presto.execution; + +import com.facebook.presto.Session; +import com.facebook.presto.execution.StateMachine.StateChangeListener; +import com.facebook.presto.server.BasicQueryInfo; +import com.facebook.presto.server.BasicQueryStats; +import com.facebook.presto.spi.ErrorCode; +import com.facebook.presto.spi.QueryId; +import com.facebook.presto.spi.memory.MemoryPoolId; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import io.airlift.units.DataSize; +import io.airlift.units.Duration; +import org.joda.time.DateTime; + +import java.net.URI; +import java.util.ArrayList; +import java.util.List; +import java.util.Optional; +import java.util.OptionalDouble; + +import static com.facebook.presto.SystemSessionProperties.QUERY_PRIORITY; +import static com.facebook.presto.execution.QueryState.FAILED; +import static com.facebook.presto.execution.QueryState.FINISHED; +import static com.facebook.presto.execution.QueryState.QUEUED; +import static com.facebook.presto.execution.QueryState.RUNNING; +import static com.facebook.presto.testing.TestingSession.testSessionBuilder; +import static io.airlift.units.DataSize.Unit.BYTE; +import static io.airlift.units.DataSize.succinctBytes; +import static java.util.concurrent.TimeUnit.MILLISECONDS; +import static java.util.concurrent.TimeUnit.NANOSECONDS; + +public class MockManagedQueryExecution + implements ManagedQueryExecution +{ + private final List> listeners = new ArrayList<>(); + private final DataSize memoryUsage; + private final Duration cpuUsage; + private final Session session; + private QueryState state = QUEUED; + private Throwable failureCause; + + public MockManagedQueryExecution(long memoryUsage) + { + this(memoryUsage, "query_id", 1); + } + + public MockManagedQueryExecution(long memoryUsage, String queryId, int priority) + { + this(memoryUsage, queryId, priority, new Duration(0, MILLISECONDS)); + } + + public MockManagedQueryExecution(long memoryUsage, String queryId, int priority, Duration cpuUsage) + { + this.memoryUsage = succinctBytes(memoryUsage); + this.cpuUsage = cpuUsage; + this.session = testSessionBuilder() + .setSystemProperty(QUERY_PRIORITY, String.valueOf(priority)) + .build(); + } + + public void complete() + { + state = FINISHED; + fireStateChange(); + } + + public Throwable getThrowable() + { + return failureCause; + } + + @Override + public Session getSession() + { + return session; + } + + @Override + public Optional getErrorCode() + { + return Optional.empty(); + } + + @Override + public BasicQueryInfo getBasicQueryInfo() + { + return new BasicQueryInfo( + new QueryId("test"), + session.toSessionRepresentation(), + Optional.empty(), + state, + new MemoryPoolId("test"), + !state.isDone(), + URI.create("http://test"), + "SELECT 1", + new BasicQueryStats( + new DateTime(1), + new DateTime(2), + new Duration(3, NANOSECONDS), + new Duration(4, NANOSECONDS), + new Duration(5, NANOSECONDS), + 6, + 7, + 8, + 9, + new DataSize(14, BYTE), + 15, + 16.0, + new DataSize(17, BYTE), + new DataSize(18, BYTE), + new DataSize(19, BYTE), + new Duration(21, NANOSECONDS), + new Duration(22, NANOSECONDS), + false, + ImmutableSet.of(), + OptionalDouble.empty()), + null, + null, + Optional.empty(), + ImmutableList.of()); + } + + @Override + public DataSize getUserMemoryReservation() + { + return memoryUsage; + } + + @Override + public DataSize getTotalMemoryReservation() + { + return memoryUsage; + } + + @Override + public Duration getTotalCpuTime() + { + return cpuUsage; + } + + public QueryState getState() + { + return state; + } + + @Override + public void startWaitingForResources() + { + state = RUNNING; + fireStateChange(); + } + + @Override + public void fail(Throwable cause) + { + state = FAILED; + failureCause = cause; + fireStateChange(); + } + + @Override + public boolean isDone() + { + return getState().isDone(); + } + + @Override + public void addStateChangeListener(StateChangeListener stateChangeListener) + { + listeners.add(stateChangeListener); + } + + private void fireStateChange() + { + for (StateChangeListener listener : listeners) { + listener.stateChanged(state); + } + } +} diff --git a/presto-main/src/test/java/com/facebook/presto/execution/MockQueryExecution.java b/presto-main/src/test/java/com/facebook/presto/execution/MockQueryExecution.java deleted file mode 100644 index d097cbfb4a4a3..0000000000000 --- a/presto-main/src/test/java/com/facebook/presto/execution/MockQueryExecution.java +++ /dev/null @@ -1,346 +0,0 @@ -/* - * 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 com.facebook.presto.execution; - -import com.facebook.presto.Session; -import com.facebook.presto.execution.StateMachine.StateChangeListener; -import com.facebook.presto.memory.VersionedMemoryPoolId; -import com.facebook.presto.server.BasicQueryInfo; -import com.facebook.presto.spi.ErrorCode; -import com.facebook.presto.spi.QueryId; -import com.facebook.presto.spi.memory.MemoryPoolId; -import com.facebook.presto.spi.resourceGroups.ResourceGroupId; -import com.facebook.presto.sql.planner.Plan; -import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableMap; -import com.google.common.collect.ImmutableSet; -import com.google.common.util.concurrent.ListenableFuture; -import io.airlift.units.DataSize; -import io.airlift.units.Duration; -import org.joda.time.DateTime; - -import java.net.URI; -import java.util.ArrayList; -import java.util.List; -import java.util.Optional; -import java.util.function.Consumer; - -import static com.facebook.presto.SystemSessionProperties.QUERY_PRIORITY; -import static com.facebook.presto.execution.QueryState.FAILED; -import static com.facebook.presto.execution.QueryState.FINISHED; -import static com.facebook.presto.execution.QueryState.QUEUED; -import static com.facebook.presto.execution.QueryState.RUNNING; -import static com.facebook.presto.testing.TestingSession.testSessionBuilder; -import static com.google.common.util.concurrent.Futures.immediateFuture; -import static io.airlift.units.DataSize.Unit.BYTE; -import static io.airlift.units.DataSize.succinctBytes; -import static java.util.concurrent.TimeUnit.MILLISECONDS; -import static java.util.concurrent.TimeUnit.NANOSECONDS; - -public class MockQueryExecution - implements QueryExecution -{ - private final List> listeners = new ArrayList<>(); - private final DataSize memoryUsage; - private final Duration cpuUsage; - private final Session session; - private final QueryId queryId; - private QueryState state = QUEUED; - private Throwable failureCause; - private Optional resourceGroupId; - - public MockQueryExecution(long memoryUsage) - { - this(memoryUsage, "query_id", 1); - } - - public MockQueryExecution(long memoryUsage, String queryId, int priority) - { - this(memoryUsage, queryId, priority, new Duration(0, MILLISECONDS)); - } - - public MockQueryExecution(long memoryUsage, String queryId, int priority, Duration cpuUsage) - { - this.memoryUsage = succinctBytes(memoryUsage); - this.cpuUsage = cpuUsage; - this.session = testSessionBuilder() - .setSystemProperty(QUERY_PRIORITY, String.valueOf(priority)) - .build(); - this.resourceGroupId = Optional.empty(); - this.queryId = new QueryId(queryId); - } - - public void complete() - { - state = FINISHED; - fireStateChange(); - } - - @Override - public QueryId getQueryId() - { - return queryId; - } - - @Override - public QueryInfo getQueryInfo() - { - return new QueryInfo( - new QueryId("test"), - session.toSessionRepresentation(), - state, - new MemoryPoolId("test"), - !state.isDone(), - URI.create("http://test"), - ImmutableList.of(), - "SELECT 1", - new QueryStats( - new DateTime(1), - new DateTime(2), - new DateTime(3), - new DateTime(4), - new Duration(6, NANOSECONDS), - new Duration(5, NANOSECONDS), - new Duration(31, NANOSECONDS), - new Duration(41, NANOSECONDS), - new Duration(7, NANOSECONDS), - new Duration(8, NANOSECONDS), - - new Duration(100, NANOSECONDS), - - 9, - 10, - 11, - - 12, - 13, - 15, - 30, - 16, - - 17.0, - new DataSize(18, BYTE), - new DataSize(19, BYTE), - new DataSize(20, BYTE), - new DataSize(21, BYTE), - new DataSize(22, BYTE), - new DataSize(23, BYTE), - - true, - new Duration(20, NANOSECONDS), - new Duration(21, NANOSECONDS), - new Duration(23, NANOSECONDS), - false, - ImmutableSet.of(), - - new DataSize(24, BYTE), - 25, - - new DataSize(26, BYTE), - 27, - - new DataSize(28, BYTE), - 29, - 30, - new DataSize(31, BYTE), - new DataSize(32, BYTE), - new DataSize(33, BYTE), - ImmutableList.of(), - ImmutableList.of()), - Optional.empty(), - Optional.empty(), - Optional.empty(), - ImmutableMap.of(), - ImmutableSet.of(), - ImmutableMap.of(), - ImmutableMap.of(), - ImmutableSet.of(), - Optional.empty(), - false, - "", - Optional.empty(), - null, - null, - ImmutableList.of(), - ImmutableSet.of(), - Optional.empty(), - state.isDone(), - Optional.empty(), - Optional.empty()); - } - - @Override - public QueryState getState() - { - return state; - } - - @Override - public Plan getQueryPlan() - { - throw new UnsupportedOperationException(); - } - - public Throwable getThrowable() - { - return failureCause; - } - - @Override - public void addOutputInfoListener(Consumer listener) - { - // no-op - } - - @Override - public ListenableFuture getStateChange(QueryState currentState) - { - return immediateFuture(state); - } - - @Override - public VersionedMemoryPoolId getMemoryPool() - { - throw new UnsupportedOperationException(); - } - - @Override - public void setMemoryPool(VersionedMemoryPoolId poolId) - { - throw new UnsupportedOperationException(); - } - - @Override - public Session getSession() - { - return session; - } - - @Override - public DateTime getCreateTime() - { - return getQueryInfo().getQueryStats().getCreateTime(); - } - - @Override - public Optional getExecutionStartTime() - { - return Optional.ofNullable(getQueryInfo().getQueryStats().getExecutionStartTime()); - } - - @Override - public DateTime getLastHeartbeat() - { - return getQueryInfo().getQueryStats().getLastHeartbeat(); - } - - @Override - public Optional getEndTime() - { - return Optional.ofNullable(getQueryInfo().getQueryStats().getEndTime()); - } - - @Override - public Optional getErrorCode() - { - return Optional.ofNullable(getQueryInfo().getFailureInfo()).map(ExecutionFailureInfo::getErrorCode); - } - - @Override - public BasicQueryInfo getBasicQueryInfo() - { - return new BasicQueryInfo(getQueryInfo()); - } - - @Override - public DataSize getUserMemoryReservation() - { - return memoryUsage; - } - - @Override - public DataSize getTotalMemoryReservation() - { - return memoryUsage; - } - - @Override - public Duration getTotalCpuTime() - { - return cpuUsage; - } - - @Override - public void start() - { - state = RUNNING; - fireStateChange(); - } - - @Override - public void fail(Throwable cause) - { - state = FAILED; - failureCause = cause; - fireStateChange(); - } - - @Override - public boolean isDone() - { - return getState().isDone(); - } - - @Override - public void cancelQuery() - { - state = FAILED; - fireStateChange(); - } - - @Override - public void cancelStage(StageId stageId) - { - throw new UnsupportedOperationException(); - } - - @Override - public void recordHeartbeat() - { - } - - @Override - public void pruneInfo() - { - } - - @Override - public void addStateChangeListener(StateChangeListener stateChangeListener) - { - listeners.add(stateChangeListener); - } - - @Override - public void addFinalQueryInfoListener(StateChangeListener stateChangeListener) - { - throw new UnsupportedOperationException(); - } - - private void fireStateChange() - { - for (StateChangeListener listener : listeners) { - listener.stateChanged(state); - } - } -} diff --git a/presto-main/src/test/java/com/facebook/presto/execution/TestQueryManagerConfig.java b/presto-main/src/test/java/com/facebook/presto/execution/TestQueryManagerConfig.java index 04d67750d51bf..5206566f78e34 100644 --- a/presto-main/src/test/java/com/facebook/presto/execution/TestQueryManagerConfig.java +++ b/presto-main/src/test/java/com/facebook/presto/execution/TestQueryManagerConfig.java @@ -49,8 +49,6 @@ public void testDefaults() .setQueryMaxRunTime(new Duration(100, TimeUnit.DAYS)) .setQueryMaxExecutionTime(new Duration(100, TimeUnit.DAYS)) .setQueryMaxCpuTime(new Duration(1_000_000_000, TimeUnit.DAYS)) - .setInitializationRequiredWorkers(1) - .setInitializationTimeout(new Duration(5, TimeUnit.MINUTES)) .setRequiredWorkers(1) .setRequiredWorkersMaxWait(new Duration(5, TimeUnit.MINUTES))); } @@ -80,8 +78,6 @@ public void testExplicitPropertyMappings() .put("query.max-run-time", "2h") .put("query.max-execution-time", "3h") .put("query.max-cpu-time", "2d") - .put("query-manager.initialization-required-workers", "200") - .put("query-manager.initialization-timeout", "1m") .put("query-manager.required-workers", "333") .put("query-manager.required-workers-max-wait", "33m") .build(); @@ -108,8 +104,6 @@ public void testExplicitPropertyMappings() .setQueryMaxRunTime(new Duration(2, TimeUnit.HOURS)) .setQueryMaxExecutionTime(new Duration(3, TimeUnit.HOURS)) .setQueryMaxCpuTime(new Duration(2, TimeUnit.DAYS)) - .setInitializationRequiredWorkers(200) - .setInitializationTimeout(new Duration(1, TimeUnit.MINUTES)) .setRequiredWorkers(333) .setRequiredWorkersMaxWait(new Duration(33, TimeUnit.MINUTES)); diff --git a/presto-main/src/test/java/com/facebook/presto/execution/TestQueryStateMachine.java b/presto-main/src/test/java/com/facebook/presto/execution/TestQueryStateMachine.java index cfb57613b04df..7518ec7630fd4 100644 --- a/presto-main/src/test/java/com/facebook/presto/execution/TestQueryStateMachine.java +++ b/presto-main/src/test/java/com/facebook/presto/execution/TestQueryStateMachine.java @@ -47,6 +47,7 @@ import java.util.function.Consumer; import static com.facebook.presto.SessionTestUtils.TEST_SESSION; +import static com.facebook.presto.execution.QueryState.DISPATCHING; import static com.facebook.presto.execution.QueryState.FAILED; import static com.facebook.presto.execution.QueryState.FINISHED; import static com.facebook.presto.execution.QueryState.FINISHING; @@ -102,6 +103,9 @@ public void testBasicStateChanges() QueryStateMachine stateMachine = createQueryStateMachine(); assertState(stateMachine, QUEUED); + assertTrue(stateMachine.transitionToDispatching()); + assertState(stateMachine, DISPATCHING); + assertTrue(stateMachine.transitionToPlanning()); assertState(stateMachine, PLANNING); @@ -125,6 +129,9 @@ public void testStateChangesWithResourceWaiting() assertTrue(stateMachine.transitionToWaitingForResources()); assertState(stateMachine, WAITING_FOR_RESOURCES); + assertTrue(stateMachine.transitionToDispatching()); + assertState(stateMachine, DISPATCHING); + assertTrue(stateMachine.transitionToPlanning()); assertState(stateMachine, PLANNING); @@ -145,6 +152,7 @@ public void testQueued() // all time before the first state transition is accounted to queueing assertAllTimeSpentInQueueing(QUEUED, queryStateMachine -> {}); assertAllTimeSpentInQueueing(WAITING_FOR_RESOURCES, QueryStateMachine::transitionToWaitingForResources); + assertAllTimeSpentInQueueing(DISPATCHING, QueryStateMachine::transitionToDispatching); assertAllTimeSpentInQueueing(PLANNING, QueryStateMachine::transitionToPlanning); assertAllTimeSpentInQueueing(STARTING, QueryStateMachine::transitionToStarting); assertAllTimeSpentInQueueing(RUNNING, QueryStateMachine::transitionToRunning); @@ -169,6 +177,7 @@ private void assertAllTimeSpentInQueueing(QueryState expectedState, Consumer queries = new TreeMap<>(); + SortedMap queries = new TreeMap<>(); Random random = new Random(); for (int i = 0; i < 100; i++) { @@ -357,7 +357,7 @@ public void testPriorityScheduling() } while (queries.containsKey(priority)); - MockQueryExecution query = new MockQueryExecution(0, "query_id", priority); + MockManagedQueryExecution query = new MockManagedQueryExecution(0, "query_id", priority); if (random.nextBoolean()) { group1.run(query); } @@ -369,10 +369,10 @@ public void testPriorityScheduling() root.setHardConcurrencyLimit(1); - List orderedQueries = new ArrayList<>(queries.values()); + List orderedQueries = new ArrayList<>(queries.values()); reverse(orderedQueries); - for (MockQueryExecution query : orderedQueries) { + for (MockManagedQueryExecution query : orderedQueries) { root.processQueuedQueries(); assertEquals(query.getState(), RUNNING); query.complete(); @@ -400,14 +400,14 @@ public void testWeightedScheduling() group2.setSoftConcurrencyLimit(2); group2.setSchedulingWeight(2); - Set group1Queries = fillGroupTo(group1, ImmutableSet.of(), 2); - Set group2Queries = fillGroupTo(group2, ImmutableSet.of(), 2); + Set group1Queries = fillGroupTo(group1, ImmutableSet.of(), 2); + Set group2Queries = fillGroupTo(group2, ImmutableSet.of(), 2); root.setHardConcurrencyLimit(1); int group2Ran = 0; for (int i = 0; i < 1000; i++) { - for (Iterator iterator = group1Queries.iterator(); iterator.hasNext(); ) { - MockQueryExecution query = iterator.next(); + for (Iterator iterator = group1Queries.iterator(); iterator.hasNext(); ) { + MockManagedQueryExecution query = iterator.next(); if (query.getState() == RUNNING) { query.complete(); iterator.remove(); @@ -452,8 +452,8 @@ public void testWeightedFairScheduling() group2.setSoftConcurrencyLimit(2); group2.setSchedulingWeight(2); - Set group1Queries = fillGroupTo(group1, ImmutableSet.of(), 4); - Set group2Queries = fillGroupTo(group2, ImmutableSet.of(), 4); + Set group1Queries = fillGroupTo(group1, ImmutableSet.of(), 4); + Set group2Queries = fillGroupTo(group2, ImmutableSet.of(), 4); root.setHardConcurrencyLimit(3); int group1Ran = 0; @@ -502,9 +502,9 @@ public void testWeightedFairSchedulingEqualWeights() group3.setSoftConcurrencyLimit(2); group3.setSchedulingWeight(2); - Set group1Queries = fillGroupTo(group1, ImmutableSet.of(), 4); - Set group2Queries = fillGroupTo(group2, ImmutableSet.of(), 4); - Set group3Queries = fillGroupTo(group3, ImmutableSet.of(), 4); + Set group1Queries = fillGroupTo(group1, ImmutableSet.of(), 4); + Set group2Queries = fillGroupTo(group2, ImmutableSet.of(), 4); + Set group3Queries = fillGroupTo(group3, ImmutableSet.of(), 4); root.setHardConcurrencyLimit(4); int group1Ran = 0; @@ -554,8 +554,8 @@ public void testWeightedFairSchedulingNoStarvation() group2.setSoftConcurrencyLimit(2); group2.setSchedulingWeight(2); - Set group1Queries = fillGroupTo(group1, ImmutableSet.of(), 4); - Set group2Queries = fillGroupTo(group2, ImmutableSet.of(), 4); + Set group1Queries = fillGroupTo(group1, ImmutableSet.of(), 4); + Set group2Queries = fillGroupTo(group2, ImmutableSet.of(), 4); root.setHardConcurrencyLimit(1); int group1Ran = 0; @@ -614,7 +614,7 @@ public void testGetInfo() rootBY.setHardConcurrencyLimit(10); // Queue 40 queries (= maxQueuedQueries (40) + maxRunningQueries (0)) - Set queries = fillGroupTo(rootAX, ImmutableSet.of(), 10, false); + Set queries = fillGroupTo(rootAX, ImmutableSet.of(), 10, false); queries.addAll(fillGroupTo(rootAY, ImmutableSet.of(), 10, false)); queries.addAll(fillGroupTo(rootBX, ImmutableSet.of(), 10, true)); queries.addAll(fillGroupTo(rootBY, ImmutableSet.of(), 10, true)); @@ -631,9 +631,9 @@ public void testGetInfo() assertEquals(info.getNumQueuedQueries(), 36); // Complete running queries - Iterator iterator = queries.iterator(); + Iterator iterator = queries.iterator(); while (iterator.hasNext()) { - MockQueryExecution query = iterator.next(); + MockManagedQueryExecution query = iterator.next(); if (query.getState() == RUNNING) { query.complete(); iterator.remove(); @@ -692,7 +692,7 @@ public void testGetResourceGroupStateInfo() rootAY.setMaxQueuedQueries(10); rootAY.setHardConcurrencyLimit(10); - Set queries = fillGroupTo(rootAX, ImmutableSet.of(), 5, false); + Set queries = fillGroupTo(rootAX, ImmutableSet.of(), 5, false); queries.addAll(fillGroupTo(rootAY, ImmutableSet.of(), 5, false)); queries.addAll(fillGroupTo(rootB, ImmutableSet.of(), 10, true)); @@ -769,7 +769,7 @@ public void testGetBlockedQueuedQueries() rootBY.setHardConcurrencyLimit(5); // Queue 40 queries (= maxQueuedQueries (40) + maxRunningQueries (0)) - Set queries = fillGroupTo(rootAX, ImmutableSet.of(), 10, false); + Set queries = fillGroupTo(rootAX, ImmutableSet.of(), 10, false); queries.addAll(fillGroupTo(rootAY, ImmutableSet.of(), 10, false)); queries.addAll(fillGroupTo(rootBX, ImmutableSet.of(), 10, true)); queries.addAll(fillGroupTo(rootBY, ImmutableSet.of(), 10, true)); @@ -793,11 +793,11 @@ public void testGetBlockedQueuedQueries() assertEquals(rootBY.getWaitingQueuedQueries(), 6); } - private static int completeGroupQueries(Set groupQueries) + private static int completeGroupQueries(Set groupQueries) { int groupRan = 0; - for (Iterator iterator = groupQueries.iterator(); iterator.hasNext(); ) { - MockQueryExecution query = iterator.next(); + for (Iterator iterator = groupQueries.iterator(); iterator.hasNext(); ) { + MockManagedQueryExecution query = iterator.next(); if (query.getState() == RUNNING) { query.complete(); iterator.remove(); @@ -807,17 +807,17 @@ private static int completeGroupQueries(Set groupQueries) return groupRan; } - private static Set fillGroupTo(InternalResourceGroup group, Set existingQueries, int count) + private static Set fillGroupTo(InternalResourceGroup group, Set existingQueries, int count) { return fillGroupTo(group, existingQueries, count, false); } - private static Set fillGroupTo(InternalResourceGroup group, Set existingQueries, int count, boolean queryPriority) + private static Set fillGroupTo(InternalResourceGroup group, Set existingQueries, int count, boolean queryPriority) { int existingCount = existingQueries.size(); - Set queries = new HashSet<>(existingQueries); + Set queries = new HashSet<>(existingQueries); for (int i = 0; i < count - existingCount; i++) { - MockQueryExecution query = new MockQueryExecution(0, group.getId().toString().replace(".", "") + Integer.toString(i), queryPriority ? i + 1 : 1); + MockManagedQueryExecution query = new MockManagedQueryExecution(0, group.getId().toString().replace(".", "") + Integer.toString(i), queryPriority ? i + 1 : 1); queries.add(query); group.run(query); } diff --git a/presto-main/src/test/java/com/facebook/presto/server/TestBasicQueryInfo.java b/presto-main/src/test/java/com/facebook/presto/server/TestBasicQueryInfo.java index be008b0ed58ac..a0d7fba28e8cf 100644 --- a/presto-main/src/test/java/com/facebook/presto/server/TestBasicQueryInfo.java +++ b/presto-main/src/test/java/com/facebook/presto/server/TestBasicQueryInfo.java @@ -60,6 +60,7 @@ public void testConstructor() Duration.valueOf("8m"), Duration.valueOf("7m"), Duration.valueOf("34m"), + Duration.valueOf("35m"), Duration.valueOf("44m"), Duration.valueOf("9m"), Duration.valueOf("10m"), diff --git a/presto-main/src/test/java/com/facebook/presto/server/TestQueryStateInfo.java b/presto-main/src/test/java/com/facebook/presto/server/TestQueryStateInfo.java index f2ed1d87a0743..8c959615c4f3a 100644 --- a/presto-main/src/test/java/com/facebook/presto/server/TestQueryStateInfo.java +++ b/presto-main/src/test/java/com/facebook/presto/server/TestQueryStateInfo.java @@ -110,6 +110,7 @@ private QueryInfo createQueryInfo(String queryId, QueryState state, String query Duration.valueOf("8m"), Duration.valueOf("7m"), Duration.valueOf("34m"), + Duration.valueOf("35m"), Duration.valueOf("44m"), Duration.valueOf("9m"), Duration.valueOf("10m"), diff --git a/presto-product-tests/src/test/resources/com/facebook/presto/tests/querystats/single_query_info_response.json b/presto-product-tests/src/test/resources/com/facebook/presto/tests/querystats/single_query_info_response.json index a8ee7298a8b08..fa375b0f02b11 100644 --- a/presto-product-tests/src/test/resources/com/facebook/presto/tests/querystats/single_query_info_response.json +++ b/presto-product-tests/src/test/resources/com/facebook/presto/tests/querystats/single_query_info_response.json @@ -29,6 +29,7 @@ "elapsedTime": "134.00ms", "resourceWaitingTime": "11.00ms", "queuedTime": "2.11ms", + "dispatchingTime": "1111.00ms", "executionTime": "13.00ms", "analysisTime": "7.47ms", "totalPlanningTime": "9.99ms", diff --git a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestDistributedQueries.java b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestDistributedQueries.java index 0e9754143f12d..90d6b1a87d229 100644 --- a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestDistributedQueries.java +++ b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestDistributedQueries.java @@ -15,6 +15,7 @@ import com.facebook.presto.Session; import com.facebook.presto.SystemSessionProperties; +import com.facebook.presto.dispatcher.DispatchManager; import com.facebook.presto.execution.QueryInfo; import com.facebook.presto.execution.QueryManager; import com.facebook.presto.server.BasicQueryInfo; @@ -817,8 +818,9 @@ public void testQueryLoggingCount() // The completed queries counter is updated in a final query info listener, which is called eventually. // Therefore, here we wait until the value of this counter gets stable. - long beforeCompletedQueriesCount = waitUntilStable(() -> queryManager.getStats().getCompletedQueries().getTotalCount(), new Duration(5, SECONDS)); - long beforeSubmittedQueriesCount = queryManager.getStats().getSubmittedQueries().getTotalCount(); + DispatchManager dispatchManager = ((DistributedQueryRunner) getQueryRunner()).getCoordinator().getDispatchManager(); + long beforeCompletedQueriesCount = waitUntilStable(() -> dispatchManager.getStats().getCompletedQueries().getTotalCount(), new Duration(5, SECONDS)); + long beforeSubmittedQueriesCount = dispatchManager.getStats().getSubmittedQueries().getTotalCount(); assertUpdate("CREATE TABLE test_query_logging_count AS SELECT 1 foo_1, 2 foo_2_4", 1); assertQuery("SELECT foo_1, foo_2_4 FROM test_query_logging_count", "SELECT 1, 2"); assertUpdate("DROP TABLE test_query_logging_count"); @@ -826,9 +828,9 @@ public void testQueryLoggingCount() // TODO: Figure out a better way of synchronization assertUntilTimeout( - () -> assertEquals(queryManager.getStats().getCompletedQueries().getTotalCount() - beforeCompletedQueriesCount, 4), + () -> assertEquals(dispatchManager.getStats().getCompletedQueries().getTotalCount() - beforeCompletedQueriesCount, 4), new Duration(1, MINUTES)); - assertEquals(queryManager.getStats().getSubmittedQueries().getTotalCount() - beforeSubmittedQueriesCount, 4); + assertEquals(dispatchManager.getStats().getSubmittedQueries().getTotalCount() - beforeSubmittedQueriesCount, 4); }); } @@ -1016,7 +1018,7 @@ public void testWrittenStats() String sql = "CREATE TABLE test_written_stats AS SELECT * FROM nation"; DistributedQueryRunner distributedQueryRunner = (DistributedQueryRunner) getQueryRunner(); ResultWithQueryId resultResultWithQueryId = distributedQueryRunner.executeWithQueryId(getSession(), sql); - QueryInfo queryInfo = distributedQueryRunner.getQueryInfo(resultResultWithQueryId.getQueryId()); + QueryInfo queryInfo = distributedQueryRunner.getCoordinator().getQueryManager().getFullQueryInfo(resultResultWithQueryId.getQueryId()); assertEquals(queryInfo.getQueryStats().getOutputPositions(), 1L); assertEquals(queryInfo.getQueryStats().getWrittenOutputPositions(), 25L); @@ -1024,7 +1026,7 @@ public void testWrittenStats() sql = "INSERT INTO test_written_stats SELECT * FROM nation LIMIT 10"; resultResultWithQueryId = distributedQueryRunner.executeWithQueryId(getSession(), sql); - queryInfo = distributedQueryRunner.getQueryInfo(resultResultWithQueryId.getQueryId()); + queryInfo = distributedQueryRunner.getCoordinator().getQueryManager().getFullQueryInfo(resultResultWithQueryId.getQueryId()); assertEquals(queryInfo.getQueryStats().getOutputPositions(), 1L); assertEquals(queryInfo.getQueryStats().getWrittenOutputPositions(), 10L); diff --git a/presto-tests/src/main/java/com/facebook/presto/tests/DistributedQueryRunner.java b/presto-tests/src/main/java/com/facebook/presto/tests/DistributedQueryRunner.java index fff505dd730df..00bb9fad9c299 100644 --- a/presto-tests/src/main/java/com/facebook/presto/tests/DistributedQueryRunner.java +++ b/presto-tests/src/main/java/com/facebook/presto/tests/DistributedQueryRunner.java @@ -17,7 +17,6 @@ import com.facebook.presto.Session.SessionBuilder; import com.facebook.presto.connector.ConnectorId; import com.facebook.presto.cost.StatsCalculator; -import com.facebook.presto.execution.QueryInfo; import com.facebook.presto.execution.QueryManager; import com.facebook.presto.execution.warnings.WarningCollector; import com.facebook.presto.metadata.AllNodes; @@ -415,11 +414,6 @@ public Plan createPlan(Session session, String sql, WarningCollector warningColl return queryPlan; } - public QueryInfo getQueryInfo(QueryId queryId) - { - return coordinator.getQueryManager().getFullQueryInfo(queryId); - } - public Plan getQueryPlan(QueryId queryId) { return coordinator.getQueryPlan(queryId); diff --git a/presto-tests/src/test/java/com/facebook/presto/execution/TestEventListener.java b/presto-tests/src/test/java/com/facebook/presto/execution/TestEventListener.java index 496870e97fed0..64d4aab2b3cb0 100644 --- a/presto-tests/src/test/java/com/facebook/presto/execution/TestEventListener.java +++ b/presto-tests/src/test/java/com/facebook/presto/execution/TestEventListener.java @@ -178,7 +178,7 @@ public void testOutputStats() MaterializedResult result = runQueryAndWaitForEvents("SELECT 1 FROM lineitem", expectedEvents); QueryCreatedEvent queryCreatedEvent = getOnlyElement(generatedEvents.getQueryCreatedEvents()); QueryCompletedEvent queryCompletedEvent = getOnlyElement(generatedEvents.getQueryCompletedEvents()); - QueryStats queryStats = queryRunner.getQueryInfo(new QueryId(queryCreatedEvent.getMetadata().getQueryId())).getQueryStats(); + QueryStats queryStats = queryRunner.getCoordinator().getQueryManager().getFullQueryInfo(new QueryId(queryCreatedEvent.getMetadata().getQueryId())).getQueryStats(); assertTrue(queryStats.getOutputDataSize().toBytes() > 0L); assertTrue(queryCompletedEvent.getStatistics().getOutputBytes() > 0L); @@ -188,7 +188,7 @@ public void testOutputStats() runQueryAndWaitForEvents("SELECT COUNT(1) FROM lineitem", expectedEvents); queryCreatedEvent = getOnlyElement(generatedEvents.getQueryCreatedEvents()); queryCompletedEvent = getOnlyElement(generatedEvents.getQueryCompletedEvents()); - queryStats = queryRunner.getQueryInfo(new QueryId(queryCreatedEvent.getMetadata().getQueryId())).getQueryStats(); + queryStats = queryRunner.getCoordinator().getQueryManager().getFullQueryInfo(new QueryId(queryCreatedEvent.getMetadata().getQueryId())).getQueryStats(); assertTrue(queryStats.getOutputDataSize().toBytes() > 0L); assertTrue(queryCompletedEvent.getStatistics().getOutputBytes() > 0L); diff --git a/presto-tests/src/test/java/com/facebook/presto/execution/TestQueryRunnerUtil.java b/presto-tests/src/test/java/com/facebook/presto/execution/TestQueryRunnerUtil.java index aae63ed6b80e1..b7d0320880ae1 100644 --- a/presto-tests/src/test/java/com/facebook/presto/execution/TestQueryRunnerUtil.java +++ b/presto-tests/src/test/java/com/facebook/presto/execution/TestQueryRunnerUtil.java @@ -14,6 +14,7 @@ package com.facebook.presto.execution; import com.facebook.presto.Session; +import com.facebook.presto.dispatcher.DispatchManager; import com.facebook.presto.server.BasicQueryInfo; import com.facebook.presto.spi.QueryId; import com.facebook.presto.tests.DistributedQueryRunner; @@ -35,14 +36,14 @@ private TestQueryRunnerUtil() {} public static QueryId createQuery(DistributedQueryRunner queryRunner, Session session, String sql) { - QueryManager queryManager = queryRunner.getCoordinator().getQueryManager(); - getFutureValue(queryManager.createQuery(session.getQueryId(), new TestingSessionContext(session), sql)); + DispatchManager dispatchManager = queryRunner.getCoordinator().getDispatchManager(); + getFutureValue(dispatchManager.createQuery(session.getQueryId(), "slug", new TestingSessionContext(session), sql)); return session.getQueryId(); } public static void cancelQuery(DistributedQueryRunner queryRunner, QueryId queryId) { - queryRunner.getCoordinator().getQueryManager().cancelQuery(queryId); + queryRunner.getCoordinator().getDispatchManager().cancelQuery(queryId); } public static void waitForQueryState(DistributedQueryRunner queryRunner, QueryId queryId, QueryState expectedQueryState) @@ -54,17 +55,17 @@ public static void waitForQueryState(DistributedQueryRunner queryRunner, QueryId public static void waitForQueryState(DistributedQueryRunner queryRunner, QueryId queryId, Set expectedQueryStates) throws InterruptedException { - QueryManager queryManager = queryRunner.getCoordinator().getQueryManager(); + DispatchManager dispatchManager = queryRunner.getCoordinator().getDispatchManager(); do { // Heartbeat all the running queries, so they don't die while we're waiting - for (BasicQueryInfo queryInfo : queryManager.getQueries()) { + for (BasicQueryInfo queryInfo : dispatchManager.getQueries()) { if (queryInfo.getState() == RUNNING) { - queryManager.recordHeartbeat(queryInfo.getQueryId()); + dispatchManager.getQueryInfo(queryInfo.getQueryId()); } } MILLISECONDS.sleep(500); } - while (!expectedQueryStates.contains(queryManager.getQueryState(queryId))); + while (!expectedQueryStates.contains(dispatchManager.getQueryInfo(queryId).getState())); } public static DistributedQueryRunner createQueryRunner() diff --git a/presto-tests/src/test/java/com/facebook/presto/execution/TestQueues.java b/presto-tests/src/test/java/com/facebook/presto/execution/TestQueues.java index ffd6383cf49e0..ed5f768683665 100644 --- a/presto-tests/src/test/java/com/facebook/presto/execution/TestQueues.java +++ b/presto-tests/src/test/java/com/facebook/presto/execution/TestQueues.java @@ -14,6 +14,7 @@ package com.facebook.presto.execution; import com.facebook.presto.Session; +import com.facebook.presto.dispatcher.DispatchManager; import com.facebook.presto.resourceGroups.ResourceGroupManagerPlugin; import com.facebook.presto.spi.QueryId; import com.facebook.presto.spi.resourceGroups.ResourceGroupId; @@ -304,8 +305,8 @@ private void testRejection() QueryId queryId = createQuery(queryRunner, newRejectionSession(), LONG_LASTING_QUERY); waitForQueryState(queryRunner, queryId, FAILED); - QueryManager queryManager = queryRunner.getCoordinator().getQueryManager(); - assertEquals(queryManager.getQueryInfo(queryId).getErrorCode(), QUERY_REJECTED.toErrorCode()); + DispatchManager dispatchManager = queryRunner.getCoordinator().getDispatchManager(); + assertEquals(dispatchManager.getQueryInfo(queryId).getErrorCode(), QUERY_REJECTED.toErrorCode()); } } diff --git a/presto-tests/src/test/java/com/facebook/presto/execution/resourceGroups/db/TestQueuesDb.java b/presto-tests/src/test/java/com/facebook/presto/execution/resourceGroups/db/TestQueuesDb.java index 031b26aba2f3b..4cc67a85e13fe 100644 --- a/presto-tests/src/test/java/com/facebook/presto/execution/resourceGroups/db/TestQueuesDb.java +++ b/presto-tests/src/test/java/com/facebook/presto/execution/resourceGroups/db/TestQueuesDb.java @@ -14,10 +14,12 @@ package com.facebook.presto.execution.resourceGroups.db; import com.facebook.presto.Session; +import com.facebook.presto.dispatcher.DispatchManager; import com.facebook.presto.execution.QueryManager; import com.facebook.presto.execution.resourceGroups.InternalResourceGroupManager; import com.facebook.presto.resourceGroups.db.DbResourceGroupConfigurationManager; import com.facebook.presto.resourceGroups.db.H2ResourceGroupsDao; +import com.facebook.presto.server.BasicQueryInfo; import com.facebook.presto.server.ResourceGroupInfo; import com.facebook.presto.spi.QueryId; import com.facebook.presto.spi.resourceGroups.ResourceGroupId; @@ -54,6 +56,7 @@ import static com.facebook.presto.execution.resourceGroups.db.H2TestUtil.waitForRunningQueryCount; import static com.facebook.presto.spi.StandardErrorCode.EXCEEDED_TIME_LIMIT; import static com.facebook.presto.spi.StandardErrorCode.INVALID_RESOURCE_GROUP; +import static com.facebook.presto.spi.StandardErrorCode.QUERY_QUEUE_FULL; import static com.facebook.presto.spi.StandardErrorCode.QUERY_REJECTED; import static com.facebook.presto.testing.TestingSession.testSessionBuilder; import static io.airlift.testing.Assertions.assertContains; @@ -197,8 +200,8 @@ public void testRejection() // Verify the query cannot be submitted QueryId queryId = createQuery(queryRunner, rejectingSession(), LONG_LASTING_QUERY); waitForQueryState(queryRunner, queryId, FAILED); - QueryManager queryManager = queryRunner.getCoordinator().getQueryManager(); - assertEquals(queryManager.getQueryInfo(queryId).getErrorCode(), QUERY_REJECTED.toErrorCode()); + DispatchManager dispatchManager = queryRunner.getCoordinator().getDispatchManager(); + assertEquals(dispatchManager.getQueryInfo(queryId).getErrorCode(), QUERY_REJECTED.toErrorCode()); int selectorCount = getSelectors(queryRunner).size(); dao.insertSelector(4, 100_000, "user.*", "(?i).*reject.*", null, null, null); dbConfigurationManager.load(); @@ -250,9 +253,9 @@ public void testSelectorPriority() QueryId secondQuery = createQuery(queryRunner, dashboardSession(), LONG_LASTING_QUERY); waitForQueryState(queryRunner, secondQuery, FAILED); - resourceGroup = queryManager.getFullQueryInfo(secondQuery).getResourceGroupId(); - assertTrue(resourceGroup.isPresent()); - assertEquals(resourceGroup.get(), createResourceGroupId("global", "user-user", "reject-all-queries")); + DispatchManager dispatchManager = queryRunner.getCoordinator().getDispatchManager(); + BasicQueryInfo basicQueryInfo = dispatchManager.getQueryInfo(secondQuery); + assertEquals(basicQueryInfo.getErrorCode(), QUERY_QUEUE_FULL.toErrorCode()); } @Test(timeOut = 60_000) @@ -290,12 +293,13 @@ public void testQueryExecutionTimeLimit() waitForQueryState(queryRunner, secondQuery, QUEUED); // after a 5s wait this query should still be QUEUED, not FAILED as the max execution time should be enforced after the query starts running Thread.sleep(5_000); - assertEquals(queryManager.getQueryState(secondQuery), QUEUED); + DispatchManager dispatchManager = queryRunner.getCoordinator().getDispatchManager(); + assertEquals(dispatchManager.getQueryInfo(secondQuery).getState(), QUEUED); // reconfigure the resource group to run the second query dao.updateResourceGroup(5, "dashboard-${USER}", "1MB", 1, null, 1, null, null, null, null, null, 3L, TEST_ENVIRONMENT); dbConfigurationManager.load(); // cancel the first one and let the second one start - queryManager.cancelQuery(firstQuery); + dispatchManager.cancelQuery(firstQuery); // wait until the second one is FAILED waitForQueryState(queryRunner, secondQuery, FAILED); } @@ -350,7 +354,7 @@ public void testNonLeafGroup() // Submit a query to a non-leaf resource group QueryId invalidResourceGroupQuery = createQuery(queryRunner, session, LONG_LASTING_QUERY); waitForQueryState(queryRunner, invalidResourceGroupQuery, FAILED); - assertEquals(queryRunner.getQueryInfo(invalidResourceGroupQuery).getErrorCode(), INVALID_RESOURCE_GROUP.toErrorCode()); + assertEquals(queryRunner.getCoordinator().getDispatchManager().getQueryInfo(invalidResourceGroupQuery).getErrorCode(), INVALID_RESOURCE_GROUP.toErrorCode()); } private void assertResourceGroupWithClientTags(Set clientTags, ResourceGroupId expectedResourceGroup) diff --git a/presto-tests/src/test/java/com/facebook/presto/tests/TestMetadataManager.java b/presto-tests/src/test/java/com/facebook/presto/tests/TestMetadataManager.java index e823b7e9ba1e0..5f35cdcf83bfa 100644 --- a/presto-tests/src/test/java/com/facebook/presto/tests/TestMetadataManager.java +++ b/presto-tests/src/test/java/com/facebook/presto/tests/TestMetadataManager.java @@ -14,10 +14,10 @@ package com.facebook.presto.tests; import com.facebook.presto.connector.MockConnectorFactory; -import com.facebook.presto.execution.QueryInfo; -import com.facebook.presto.execution.QueryManager; +import com.facebook.presto.dispatcher.DispatchManager; import com.facebook.presto.execution.TestingSessionContext; import com.facebook.presto.metadata.MetadataManager; +import com.facebook.presto.server.BasicQueryInfo; import com.facebook.presto.spi.Plugin; import com.facebook.presto.spi.QueryId; import com.facebook.presto.spi.connector.ConnectorFactory; @@ -113,20 +113,21 @@ public void testMetadataIsClearedAfterQueryFailed() public void testMetadataIsClearedAfterQueryCanceled() throws Exception { - QueryManager queryManager = queryRunner.getCoordinator().getQueryManager(); - QueryId queryId = queryManager.createQueryId(); - queryManager.createQuery( + DispatchManager dispatchManager = queryRunner.getCoordinator().getDispatchManager(); + QueryId queryId = dispatchManager.createQueryId(); + dispatchManager.createQuery( queryId, + "slug", new TestingSessionContext(TEST_SESSION), "SELECT * FROM lineitem") .get(); // wait until query starts running while (true) { - QueryInfo queryInfo = queryManager.getFullQueryInfo(queryId); + BasicQueryInfo queryInfo = dispatchManager.getQueryInfo(queryId); if (queryInfo.getState().isDone()) { assertEquals(queryInfo.getState(), FAILED); - throw queryInfo.getFailureInfo().toException(); + throw dispatchManager.getDispatchInfo(queryId).get().getFailureInfo().get().toException(); } if (queryInfo.getState() == RUNNING) { break; @@ -135,7 +136,7 @@ public void testMetadataIsClearedAfterQueryCanceled() } // cancel query - queryManager.cancelQuery(queryId); + dispatchManager.cancelQuery(queryId); assertEquals(metadataManager.getCatalogsByQueryId().size(), 0); } diff --git a/presto-tests/src/test/java/com/facebook/presto/tests/TestMinWorkerRequirement.java b/presto-tests/src/test/java/com/facebook/presto/tests/TestMinWorkerRequirement.java index 2997f561b3b6a..7eb819e677e09 100644 --- a/presto-tests/src/test/java/com/facebook/presto/tests/TestMinWorkerRequirement.java +++ b/presto-tests/src/test/java/com/facebook/presto/tests/TestMinWorkerRequirement.java @@ -24,67 +24,6 @@ @Test(singleThreaded = true) public class TestMinWorkerRequirement { - @Test(expectedExceptions = RuntimeException.class, expectedExceptionsMessageRegExp = "Cluster is still initializing, there are insufficient active worker nodes \\(4\\) to run query") - public void testInsufficientInitialWorkerNodes() - throws Exception - { - try (DistributedQueryRunner queryRunner = TpchQueryRunnerBuilder.builder() - .setSingleCoordinatorProperty("query-manager.initialization-required-workers", "5") - .setNodeCount(4) - .build()) { - queryRunner.execute("SELECT 1"); - fail("Expected exception due to insufficient active worker nodes"); - } - } - - @Test(expectedExceptions = RuntimeException.class, expectedExceptionsMessageRegExp = "Cluster is still initializing, there are insufficient active worker nodes \\(3\\) to run query") - public void testInsufficientInitialWorkerNodesWithCoordinatorExcluded() - throws Exception - { - try (DistributedQueryRunner queryRunner = TpchQueryRunnerBuilder.builder() - .setSingleExtraProperty("node-scheduler.include-coordinator", "false") - .setSingleCoordinatorProperty("query-manager.initialization-required-workers", "4") - .setNodeCount(4) - .build()) { - queryRunner.execute("SELECT 1"); - fail("Expected exception due to insufficient active worker nodes"); - } - } - - @Test - public void testSufficientInitialWorkerNodes() - throws Exception - { - try (DistributedQueryRunner queryRunner = TpchQueryRunnerBuilder.builder() - .setSingleCoordinatorProperty("query-manager.initialization-required-workers", "4") - .setNodeCount(4) - .build()) { - queryRunner.execute("SELECT 1"); - assertEquals(queryRunner.getCoordinator().refreshNodes().getActiveNodes().size(), 4); - - // Query should still be allowed to run if active workers drop down below the minimum required nodes - queryRunner.getServers().get(0).close(); - assertEquals(queryRunner.getCoordinator().refreshNodes().getActiveNodes().size(), 3); - queryRunner.execute("SELECT 1"); - } - } - - @Test - public void testInitializationTimeout() - throws Exception - { - try (DistributedQueryRunner queryRunner = TpchQueryRunnerBuilder.builder() - .setCoordinatorProperties(ImmutableMap.builder() - .put("query-manager.initialization-required-workers", "5") - .put("query-manager.initialization-timeout", "1ns") - .build()) - .setNodeCount(4) - .build()) { - queryRunner.execute("SELECT 1"); - assertEquals(queryRunner.getCoordinator().refreshNodes().getActiveNodes().size(), 4); - } - } - @Test(expectedExceptions = RuntimeException.class, expectedExceptionsMessageRegExp = "Insufficient active worker nodes. Waited 1.00ns for at least 5 workers, but only 4 workers are active") public void testInsufficientWorkerNodes() throws Exception diff --git a/presto-tests/src/test/java/com/facebook/presto/tests/TestQueryManager.java b/presto-tests/src/test/java/com/facebook/presto/tests/TestQueryManager.java index f4f6d82e1d771..fdc6230997696 100644 --- a/presto-tests/src/test/java/com/facebook/presto/tests/TestQueryManager.java +++ b/presto-tests/src/test/java/com/facebook/presto/tests/TestQueryManager.java @@ -13,6 +13,7 @@ */ package com.facebook.presto.tests; +import com.facebook.presto.dispatcher.DispatchManager; import com.facebook.presto.execution.QueryInfo; import com.facebook.presto.execution.QueryManager; import com.facebook.presto.execution.QueryState; @@ -60,17 +61,18 @@ public void tearDown() public void testFailQuery() throws Exception { - QueryManager queryManager = queryRunner.getCoordinator().getQueryManager(); - QueryId queryId = queryManager.createQueryId(); - queryManager.createQuery( + DispatchManager dispatchManager = queryRunner.getCoordinator().getDispatchManager(); + QueryId queryId = dispatchManager.createQueryId(); + dispatchManager.createQuery( queryId, + "slug", new TestingSessionContext(TEST_SESSION), "SELECT * FROM lineitem") .get(); // wait until query starts running while (true) { - QueryState state = queryManager.getQueryState(queryId); + QueryState state = dispatchManager.getQueryInfo(queryId).getState(); if (state.isDone()) { fail("unexpected query state: " + state); } @@ -81,6 +83,7 @@ public void testFailQuery() } // cancel query + QueryManager queryManager = queryRunner.getCoordinator().getQueryManager(); queryManager.failQuery(queryId, new PrestoException(GENERIC_INTERNAL_ERROR, "mock exception")); QueryInfo queryInfo = queryManager.getFullQueryInfo(queryId); assertEquals(queryInfo.getState(), FAILED);