TaskInfoResource.java

/*
 * 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;

import com.facebook.presto.dispatcher.DispatchManager;
import com.facebook.presto.execution.QueryManager;
import com.facebook.presto.execution.StageInfo;
import com.facebook.presto.execution.TaskId;
import com.facebook.presto.execution.TaskInfo;
import com.facebook.presto.metadata.InternalNode;
import com.facebook.presto.metadata.InternalNodeManager;
import com.facebook.presto.resourcemanager.ResourceManagerProxy;
import com.facebook.presto.spi.QueryId;
import com.google.inject.Inject;

import javax.annotation.security.RolesAllowed;
import javax.servlet.http.HttpServletRequest;
import javax.ws.rs.DefaultValue;
import javax.ws.rs.GET;
import javax.ws.rs.NotFoundException;
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.MediaType;
import javax.ws.rs.core.Response;
import javax.ws.rs.core.UriBuilder;
import javax.ws.rs.core.UriInfo;

import java.net.URI;
import java.net.UnknownHostException;
import java.util.Optional;

import static com.google.common.base.Preconditions.checkState;
import static java.util.Objects.requireNonNull;
import static javax.ws.rs.core.Response.Status.NOT_FOUND;

@Path("v1/taskInfo")
@RolesAllowed("ADMIN")
public class TaskInfoResource
{
    public static final String INCLUDE_LOCAL_QUERY_ONLY = "includeLocalQueryOnly";
    private final DispatchManager dispatchManager;
    private final QueryManager queryManager;
    private final boolean resourceManagerEnabled;
    private final InternalNodeManager internalNodeManager;
    private final Optional<ResourceManagerProxy> proxyHelper;

    @Inject
    public TaskInfoResource(
            DispatchManager dispatchManager,
            QueryManager queryManager,
            InternalNodeManager internalNodeManager,
            ServerConfig serverConfig,
            Optional<ResourceManagerProxy> proxyHelper)
    {
        this.dispatchManager = requireNonNull(dispatchManager, "dispatchManager is null");
        this.queryManager = requireNonNull(queryManager, "queryManager is null");
        this.internalNodeManager = requireNonNull(internalNodeManager, "internalNodeManager is null");
        this.resourceManagerEnabled = requireNonNull(serverConfig, "serverConfig is null").isResourceManagerEnabled();
        this.proxyHelper = requireNonNull(proxyHelper, "proxyHelper is null");
    }

    @GET
    @Path("{taskId}")
    @Produces(MediaType.APPLICATION_JSON)
    public void getTaskInfo(@PathParam("taskId") TaskId taskId,
            @QueryParam(INCLUDE_LOCAL_QUERY_ONLY) @DefaultValue("false") boolean includeLocalQueryOnly,
            @Context UriInfo uriInfo,
            @Context HttpServletRequest servletRequest,
            @Suspended AsyncResponse asyncResponse)
            throws WebApplicationException
    {
        if (requestNeedsToBeProxied(taskId, includeLocalQueryOnly)) {
            proxyTaskInfo(servletRequest, asyncResponse, uriInfo);
        }
        else {
            try {
                asyncResponse.resume(Response.ok(getTaskInfo(taskId)).build());
            }
            catch (NotFoundException e) {
                asyncResponse.resume(Response.status(NOT_FOUND).entity("Could not find the requested taskInfo").build());
            }
            catch (Exception e) {
                asyncResponse.resume(Response.serverError().entity(e.getMessage()).build());
            }
        }
    }

    private TaskInfo getTaskInfo(TaskId taskId)
    {
        QueryId queryId = taskId.getQueryId();
        try {
            Optional<StageInfo> stageInfo = queryManager.getFullQueryInfo(queryId).getOutputStage();

            if (stageInfo.isPresent()) {
                Optional<StageInfo> stage = stageInfo.get().getStageWithStageId(taskId.getStageExecutionId().getStageId());
                if (stage.isPresent()) {
                    Optional<TaskInfo> taskInfo = stage.get().getLatestAttemptExecutionInfo().getTasks().stream()
                            .filter(info -> info.getTaskId().equals(taskId))
                            .findFirst();

                    if (taskInfo.isPresent()) {
                        return taskInfo.get();
                    }
                }
            }
            throw new NotFoundException("TaskInfo not found for task id " + taskId.toString());
        }
        catch (Exception e) {
            throw new NotFoundException(e);
        }
    }

    private boolean requestNeedsToBeProxied(TaskId taskId, boolean includeLocalQueryOnly)
    {
        return !includeLocalQueryOnly
                && resourceManagerEnabled
                && !dispatchManager.isQueryPresent(taskId.getQueryId());
    }

    private URI createTaskInfoUri(UriInfo uriInfo, InternalNode resourceManagerNode)
            throws UnknownHostException
    {
        return UriBuilder.fromUri(uriInfo.getRequestUri())
                .queryParam(INCLUDE_LOCAL_QUERY_ONLY, true)
                .scheme(resourceManagerNode.getInternalUri().getScheme())
                .host(resourceManagerNode.getHostAndPort().toInetAddress().getHostName())
                .port(resourceManagerNode.getInternalUri().getPort())
                .build();
    }

    private void proxyTaskInfo(HttpServletRequest servletRequest, AsyncResponse asyncResponse, UriInfo uriInfo)
    {
        try {
            checkState(proxyHelper.isPresent());
            Optional<InternalNode> resourceManager = internalNodeManager.getResourceManagers().stream()
                    .findAny();
            if (resourceManager.isPresent()) {
                InternalNode resourceManagerNode = resourceManager.get();
                URI uri = createTaskInfoUri(uriInfo, resourceManagerNode);
                proxyHelper.get().performRequest(servletRequest, asyncResponse, uri);
            }
            else {
                asyncResponse.resume(Response.serverError().entity("Could not find the resource manager").build());
            }
        }
        catch (Exception e) {
            asyncResponse.resume(Response.serverError().entity(e.getMessage()).build());
        }
    }
}