org.apache.lens.server.query.QueryExecutionServiceImpl.java Source code

Java tutorial

Introduction

Here is the source code for org.apache.lens.server.query.QueryExecutionServiceImpl.java

Source

/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 org.apache.lens.server.query;

import static org.apache.lens.api.query.QueryStatus.Status.*;
import static org.apache.lens.server.api.LensConfConstants.*;
import static org.apache.lens.server.api.util.LensUtil.getImplementations;
import static org.apache.lens.server.error.LensServerErrorCode.SERVER_OVERLOADED;
import static org.apache.lens.server.session.LensSessionImpl.ResourceEntry;

import java.io.*;
import java.net.URI;
import java.net.URISyntaxException;
import java.sql.SQLException;
import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.ReentrantLock;

import javax.ws.rs.BadRequestException;
import javax.ws.rs.NotFoundException;
import javax.ws.rs.core.MediaType;
import javax.ws.rs.core.Response;
import javax.ws.rs.core.StreamingOutput;

import org.apache.lens.api.LensConf;
import org.apache.lens.api.LensSessionHandle;
import org.apache.lens.api.Priority;
import org.apache.lens.api.error.ErrorCollection;
import org.apache.lens.api.query.*;
import org.apache.lens.api.query.QueryStatus.Status;
import org.apache.lens.cube.metadata.DateUtil;
import org.apache.lens.driver.hive.HiveDriver;
import org.apache.lens.server.BaseLensService;
import org.apache.lens.server.LensServerConf;
import org.apache.lens.server.LensServices;
import org.apache.lens.server.api.LensConfConstants;
import org.apache.lens.server.api.driver.*;
import org.apache.lens.server.api.error.LensException;
import org.apache.lens.server.api.error.LensMultiCauseException;
import org.apache.lens.server.api.events.LensEventListener;
import org.apache.lens.server.api.health.HealthStatus;
import org.apache.lens.server.api.metrics.MethodMetricsContext;
import org.apache.lens.server.api.metrics.MethodMetricsFactory;
import org.apache.lens.server.api.metrics.MetricsService;
import org.apache.lens.server.api.query.*;
import org.apache.lens.server.api.query.collect.WaitingQueriesSelectionPolicy;
import org.apache.lens.server.api.query.comparators.*;
import org.apache.lens.server.api.query.constraint.QueryLaunchingConstraint;
import org.apache.lens.server.api.query.cost.QueryCost;
import org.apache.lens.server.api.query.events.*;
import org.apache.lens.server.api.retry.BackOffRetryHandler;
import org.apache.lens.server.api.retry.ChainedRetryPolicyDecider;
import org.apache.lens.server.api.retry.OperationRetryHandlerFactory;
import org.apache.lens.server.api.retry.RetryPolicyDecider;
import org.apache.lens.server.api.util.LensUtil;
import org.apache.lens.server.model.LogSegregationContext;
import org.apache.lens.server.model.MappedDiagnosticLogSegregationContext;
import org.apache.lens.server.query.collect.*;
import org.apache.lens.server.query.constraint.DefaultQueryLaunchingConstraintsChecker;
import org.apache.lens.server.query.constraint.QueryLaunchingConstraintsChecker;
import org.apache.lens.server.rewrite.RewriteUtil;
import org.apache.lens.server.rewrite.UserQueryToCubeQueryRewriter;
import org.apache.lens.server.session.LensSessionImpl;
import org.apache.lens.server.stats.StatisticsService;
import org.apache.lens.server.util.FairPriorityBlockingQueue;
import org.apache.lens.server.util.UtilityMethods;

import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.concurrent.BasicThreadFactory;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FSDataInputStream;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.hive.conf.HiveConf;
import org.apache.hive.service.cli.CLIService;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Lists;
import com.google.common.collect.Sets;
import lombok.*;
import lombok.extern.slf4j.Slf4j;

/**
 * The Class QueryExecutionServiceImpl.
 */
@Slf4j
public class QueryExecutionServiceImpl extends BaseLensService implements QueryExecutionService {

    /**
     * The Constant PREPARED_QUERIES_COUNTER.
     */
    public static final String PREPARED_QUERIES_COUNTER = "prepared-queries";

    /**
     * The Constant QUERY_SUBMITTER_COUNTER.
     */
    public static final String QUERY_SUBMITTER_COUNTER = "query-submitter-errors";

    /**
     * The Constant STATUS_UPDATE_COUNTER.
     */
    public static final String STATUS_UPDATE_COUNTER = "status-update-errors";

    /**
     * The Constant QUERY_PURGER_COUNTER.
     */
    public static final String QUERY_PURGER_COUNTER = "query-purger-errors";

    public static final String QUERY_EXPIRY_FAILURE_COUNTER = "query-expiry-errors";

    public static final String TOTAL_QUERIES_EXPIRED = "total-expired-queries";

    /**
     * The Constant PREPARED_QUERY_PURGER_COUNTER.
     */
    public static final String PREPARED_QUERY_PURGER_COUNTER = "prepared-query-purger-errors";

    /**
     * The millis in week.
     */
    private static long millisInWeek = 7 * 24 * 60 * 60 * 1000;

    /**
     * The Constant NAME.
     */
    public static final String NAME = "query";

    /**
     * The accepted queries.
     */
    private FairPriorityBlockingQueue<QueryContext> queuedQueries;

    /**
     * The launched queries.
     */
    @Getter
    private EstimatedQueryCollection launchedQueries;

    /**
     * The waiting queries.
     */
    private EstimatedQueryCollection waitingQueries;

    /**
     * The finished queries.
     */
    ConcurrentLinkedQueue<FinishedQuery> finishedQueries = new ConcurrentLinkedQueue<>();

    /**
     * The prepared query queue.
     */
    private DelayQueue<PreparedQueryContext> preparedQueryQueue = new DelayQueue<>();

    /**
     * The prepared queries.
     */
    private Map<QueryPrepareHandle, PreparedQueryContext> preparedQueries = new HashMap<>();

    /**
     * The all queries.
     */
    protected final ConcurrentMap<QueryHandle, QueryContext> allQueries = new ConcurrentHashMap<>();

    /**
     * The conf.
     */
    @VisibleForTesting
    Configuration conf;

    /**
     * Checks if repeated query is allowed
     */
    private boolean isDuplicateQueryAllowed;

    /**
     * The query submitter runnable.
     */
    private QuerySubmitter querySubmitterRunnable;

    /**
     * The query submitter.
     */
    protected Thread querySubmitter;

    private final AsyncStatusUpdater asyncStatusUpdater = new AsyncStatusUpdater();
    /**
     * The status poller.
     */
    private final Thread statusPoller = new Thread(new StatusPoller(), "StatusPoller");

    /**
     * The query purger.
     */
    private final Thread queryPurger = new Thread(new QueryPurger(), "QueryPurger");

    /**
     * The query expiry thread which cancels timedout queries.
     */
    private ScheduledExecutorService queryExpirer;

    /**
     * The prepare query purger.
     */
    private final Thread prepareQueryPurger = new Thread(new PreparedQueryPurger(), "PrepareQueryPurger");

    /**
     * The query result purger
     */
    private QueryResultPurger queryResultPurger;

    /**
     * The query acceptors.
     */
    private List<QueryAcceptor> queryAcceptors = new ArrayList<QueryAcceptor>();

    /**
     * The drivers.
     */
    private final Map<String, LensDriver> drivers = new HashMap<String, LensDriver>();

    /**
     * The driver selector.
     */
    private DriverSelector driverSelector;

    /**
     *  The query comparator
     */
    private Comparator<QueryContext> queryComparator;
    /**
     * The result sets.
     */
    private Map<QueryHandle, LensResultSet> resultSets = new HashMap<QueryHandle, LensResultSet>();

    /**
     * The metrics service.
     */
    private MetricsService metricsService;

    /**
     * The statistics service.
     */
    private StatisticsService statisticsService;

    /**
     * The max finished queries.
     */
    int purgeInterval;

    /**
     * The lens server dao.
     */
    LensServerDAO lensServerDao;

    /**
     * Thread pool used for running query estimates in parallel
     */
    private ExecutorService estimatePool;

    /**
     * Thread pool used for launching queries in parallel.
     */
    private ExecutorService queryLauncherPool;

    /**
     * The pool used for cancelling timed out queries.
     */
    private ExecutorService queryCancellationPool;

    private final LogSegregationContext logSegregationContext;

    private final ErrorCollection errorCollection = LensServices.get().getErrorCollection();

    private QueryLaunchingConstraintsChecker queryConstraintsChecker;

    private WaitingQueriesSelector waitingQueriesSelector;

    /**
     * If the query taken out of queued queries is added to waiting, it will be re processed from waiting
     * queries, when the next launched query is finished and removed from launched queries. If a query is
     * removed from launched queries, while query not allowed to launch was still to be added to waiting
     * queries, then waiting query will have to wait for next launched query to finish.
     * This sort of delay in waiting query execution can be avoided if removal of a query from launched
     * queries is locked using removalFromLaunchedQueriesLock, until the decision to add to waiting queries
     * and actual addition to waiting query is complete.
     * */
    private final ReentrantLock removalFromLaunchedQueriesLock = new ReentrantLock();

    private final ScheduledExecutorService waitingQueriesSelectionSvc = Executors
            .newSingleThreadScheduledExecutor();

    /**
     * This is the TTL millis for all result sets of type {@link org.apache.lens.server.api.driver.InMemoryResultSet}
     * Note : this field is non final and has a Getter and Setter for test cases
     */
    @Getter
    @Setter
    private long inMemoryResultsetTTLMillis;

    /**
     * The driver event listener.
     */
    final LensEventListener<DriverEvent> driverEventListener = new LensEventListener<DriverEvent>() {
        @Override
        public void onEvent(DriverEvent event) {
            // Need to restore session only in case of hive driver
            if (event instanceof DriverSessionStarted) {
                log.info("New driver event by driver {}", event.getDriver());
                handleDriverSessionStart(event);
            }
        }
    };
    private UserQueryToCubeQueryRewriter userQueryToCubeQueryRewriter;

    // Exponential backoff retry handler for status updates
    private BackOffRetryHandler<StatusUpdateFailureContext> statusUpdateRetryHandler;
    private RetryPolicyDecider<QueryContext> queryRetryPolicyDecider;

    /**
     * Instantiates a new query execution service impl.
     *
     * @param cliService the cli service
     * @throws LensException the lens exception
     */
    public QueryExecutionServiceImpl(CLIService cliService) throws LensException {
        super(NAME, cliService);
        this.logSegregationContext = new MappedDiagnosticLogSegregationContext();
    }

    /**
     * Initialize query acceptors and listeners.
     */
    private void initializeQueryAcceptors() throws LensException {
        String[] acceptorClasses = conf.getStrings(ACCEPTOR_CLASSES);
        if (acceptorClasses != null) {
            for (String acceptorClass : acceptorClasses) {
                try {
                    Class<?> clazz = Class.forName(acceptorClass);
                    QueryAcceptor acceptor = (QueryAcceptor) clazz.newInstance();
                    log.info("initialized query acceptor: {}", acceptor);
                    queryAcceptors.add(acceptor);
                } catch (Exception e) {
                    log.warn("Could not load the acceptor: {}", acceptorClass, e);
                    throw new LensException("Could not load acceptor" + acceptorClass, e);
                }
            }
        }
    }

    private void initializeListeners() {
        if (conf.getBoolean(QUERY_STATE_LOGGER_ENABLED, true)) {
            getEventService().addListenerForType(new QueryStatusLogger(), StatusChange.class);
            log.info("Registered query state logger");
        }
        // Add result formatter
        getEventService().addListenerForType(new ResultFormatter(this, this.logSegregationContext),
                QueryExecuted.class);
        getEventService().addListenerForType(new QueryExecutionStatisticsGenerator(getEventService()),
                QueryEnded.class);
        getEventService().addListenerForType(
                new QueryEndNotifier(this, getCliService().getHiveConf(), this.logSegregationContext),
                QueryEnded.class);
        getEventService().addListenerForType(
                new QueryEndHttpNotifier(getCliService().getHiveConf(), this.logSegregationContext),
                QueryEnded.class);
        log.info("Registered query result formatter");
    }

    /**
     * Load drivers and selector.
     *
     * @throws LensException the lens exception
     */
    private void loadDriverSelector() throws LensException {
        //Load configured Driver Selector
        try {
            Class<? extends DriverSelector> driverSelectorClass = conf.getClass(DRIVER_SELECTOR_CLASS,
                    MinQueryCostSelector.class, DriverSelector.class);
            log.info("Using driver selector class: {}", driverSelectorClass.getCanonicalName());
            driverSelector = driverSelectorClass.newInstance();
        } catch (Exception e) {
            throw new LensException(
                    "Couldn't instantiate driver selector class. Class name: " + conf.get(DRIVER_SELECTOR_CLASS)
                            + ". Please supply a valid value for " + DRIVER_SELECTOR_CLASS);
        }
    }

    private void loadQueryComparator() throws LensException {
        try {
            Class<?>[] classes = conf.getClasses(QUERY_COMPARATOR_CLASSES, MoreRetriesFirstComparator.class,
                    QueryPriorityComparator.class, FIFOQueryComparator.class);
            List<Comparator<QueryContext>> comparators = Lists.newArrayList();
            for (Class<?> clazz : classes) {
                comparators.add(clazz.asSubclass(QueryComparator.class).newInstance());
            }
            queryComparator = new ChainedComparator<>(comparators);
        } catch (Exception e) {
            throw new LensException(
                    "Couldn't instantiate query comparator class. Classes: " + conf.get(QUERY_COMPARATOR_CLASSES)
                            + ". Please supply a valid value for " + QUERY_COMPARATOR_CLASSES);
        }
    }

    /**
     * Loads drivers for the configured Driver types in lens-site.xml
     *
     * The driver's resources (<driver_type>-site.xml and other files) should be present
     * under directory conf/drivers/<driver-type>/<driver-name>
     * Example :conf/drivers/hive/h1, conf/drivers/hive/h2, conf/drivers/jdbc/mysql1, conf/drivers/jdbc/vertica1
     *
     * @throws LensException
     */
    private void loadDrivers() throws LensException {
        Collection<String> driverTypes = conf.getStringCollection(DRIVER_TYPES_AND_CLASSES);
        if (driverTypes.isEmpty()) {
            throw new LensException("No drivers configured");
        }
        File driversBaseDir = new File(
                System.getProperty(LensConfConstants.CONFIG_LOCATION, LensConfConstants.DEFAULT_CONFIG_LOCATION),
                LensConfConstants.DRIVERS_BASE_DIR);
        if (!driversBaseDir.isDirectory()) {
            throw new LensException("No drivers found at location " + driversBaseDir.getAbsolutePath());
        }
        for (String driverType : driverTypes) {
            if (StringUtils.isBlank(driverType)) {
                throw new LensException(
                        "Driver type Configuration not specified correctly. Encountered blank driver type");
            }
            String[] driverTypeAndClass = StringUtils.split(driverType.trim(), ':');
            if (driverTypeAndClass.length != 2) {
                throw new LensException("Driver type Configuration not specified correctly : " + driverType);
            }
            loadDriversForType(driverTypeAndClass[0], driverTypeAndClass[1], driversBaseDir);
        }
        if (drivers.isEmpty()) {
            throw new LensException("No drivers loaded. Please check the drivers in :" + driversBaseDir);
        }
    }

    /**
     * Loads drivers of a particular type
     *
     * @param driverType : type of driver (hive, jdbc, el, etc)
     * @param driverTypeClassName :driver class name
     * @param driversBaseDir :path for drivers directory where all driver relates resources are avilable
     * @throws LensException
     */
    private void loadDriversForType(String driverType, String driverTypeClassName, File driversBaseDir)
            throws LensException {
        File driverTypeBaseDir = new File(driversBaseDir, driverType);
        File[] driverPaths = driverTypeBaseDir.listFiles();
        if (!driverTypeBaseDir.isDirectory() || driverPaths == null || driverPaths.length == 0) {
            // May be the deployment does not have drivers of this type. We can log and ignore.
            log.warn("No drivers of type {} found in {}.", driverType, driverTypeBaseDir.getAbsolutePath());
            return;
        }
        Class driverTypeClass = null;
        try {
            driverTypeClass = conf.getClassByName(driverTypeClassName);
        } catch (Exception e) {
            log.error("Could not load the driver type class {}", driverTypeClassName, e);
            throw new LensException("Could not load Driver type class " + driverTypeClassName);
        }
        LensDriver driver = null;
        String driverName = null;
        for (File driverPath : driverPaths) {
            try {
                if (!driverPath.isDirectory()) {
                    log.warn("Ignoring resource {} while loading drivers. A driver directory was expected instead",
                            driverPath.getAbsolutePath());
                    continue;
                }
                driverName = driverPath.getName();
                driver = (LensDriver) driverTypeClass.newInstance();
                driver.configure(LensServerConf.getConfForDrivers(), driverType, driverName);
                // Register listener for all drivers. Drivers can choose to ignore this registration. As of now only Hive
                // Driver supports driver event listeners.
                driver.registerDriverEventListener(driverEventListener);
                drivers.put(driver.getFullyQualifiedName(), driver);
                log.info("Driver {} for type {} is loaded", driverPath.getName(), driverType);
            } catch (Exception e) {
                log.error("Could not load driver {} of type {}", driverPath.getName(), driverType, e);
                throw new LensException("Could not load driver " + driverPath.getName() + " of type " + driverType,
                        e);
            }
        }
    }

    private MetricsService getMetrics() {
        if (metricsService == null) {
            metricsService = LensServices.get().getService(MetricsService.NAME);
            if (metricsService == null) {
                throw new NullPointerException("Could not get metrics service");
            }
        }
        return metricsService;
    }

    private StatisticsService getStatisticsService() {
        if (statisticsService == null) {
            statisticsService = LensServices.get().getService(StatisticsService.STATS_SVC_NAME);
            if (statisticsService == null) {
                throw new NullPointerException("Could not get statistics service");
            }
        }
        return statisticsService;
    }

    /**
     * Incr counter.
     *
     * @param counter the counter
     */
    private void incrCounter(String counter) {
        getMetrics().incrCounter(QueryExecutionService.class, counter);
    }

    /**
     * Decr counter.
     *
     * @param counter the counter
     */
    private void decrCounter(String counter) {
        getMetrics().decrCounter(QueryExecutionService.class, counter);
    }

    /**
     * The Class QueryStatusLogger.
     */
    public static class QueryStatusLogger implements LensEventListener<StatusChange> {

        /**
         * The Constant STATUS_LOG.
         */
        public static final Logger STATUS_LOG = LoggerFactory.getLogger(QueryStatusLogger.class);

        /*
         * (non-Javadoc)
         *
         * @see org.apache.lens.server.api.events.LensEventListener#onEvent(org.apache.lens.server.api.events.LensEvent)
         */
        @Override
        public void onEvent(StatusChange event) throws LensException {
            STATUS_LOG.info(event.toString());
        }
    }

    /**
     * The Class FinishedQuery.
     */
    @ToString
    public class FinishedQuery {

        /**
         * The ctx.
         */
        @Getter
        private final QueryContext ctx;

        /**
         * The finish time.
         */
        private final Date finishTime;

        private LensResultSet driverRS;

        /**
         * Instantiates a new finished query.
         *
         * @param ctx the ctx
         */
        FinishedQuery(QueryContext ctx) {
            this.ctx = ctx;
            if (ctx.getEndTime() == 0) {
                this.finishTime = new Date();
                ctx.setEndTime(this.finishTime.getTime());
            } else {
                this.finishTime = new Date(ctx.getEndTime());
            }
            if (ctx.isResultAvailableInDriver()) {
                try {
                    driverRS = ctx.getSelectedDriver().fetchResultSet(getCtx());
                } catch (Exception e) {
                    log.error(
                            "Error while getting result set form driver {}. Driver result set based purging logic will be ignored",
                            ctx.getSelectedDriver(), e);
                }
            }
        }

        public boolean canBePurged() {
            try {
                if (getCtx().getStatus().getStatus().equals(SUCCESSFUL)
                        && getCtx().getStatus().isResultSetAvailable()) {
                    LensResultSet serverRS = getResultset();
                    log.info("Server Resultset for {} is {}", getQueryHandle(),
                            serverRS.getClass().getSimpleName());
                    // driverRS and serverRS will not match when server persistence is enabled. Check for purgability of both
                    // result sets in this case
                    if (driverRS != null && driverRS != serverRS) {
                        log.info("Driver Resultset for {} is {}", getQueryHandle(),
                                driverRS.getClass().getSimpleName());
                        return serverRS.canBePurged()
                                && (driverRS.canBePurged() || hasResultSetExceededTTL(driverRS));
                    } else {
                        return serverRS.canBePurged() || hasResultSetExceededTTL(serverRS);
                    }
                }
                return true;
            } catch (Throwable e) {
                log.error("Error while accessing result set for query handle while purging: {}."
                        + " Hence, going ahead with purge", getQueryHandle(), e);
                return true;
            }
        }

        /**
         * Checks the TTL for ResultSet. TTL is applicable to In Memory ResultSets only.
         *
         * @param resultSet
         * @return
         */
        private boolean hasResultSetExceededTTL(LensResultSet resultSet) {
            if (resultSet instanceof InMemoryResultSet
                    && System.currentTimeMillis() > ((InMemoryResultSet) resultSet).getCreationTime()
                            + inMemoryResultsetTTLMillis) {
                log.info("InMemoryResultSet for query {} has exceeded its TTL and is eligible for purging now",
                        getQueryHandle());
                return true;
            }
            return false;
        }

        private LensResultSet getResultset() throws LensException {
            return QueryExecutionServiceImpl.this.getResultset(getQueryHandle());
        }

        /**
         * @return the finishTime
         */
        public Date getFinishTime() {
            return finishTime;
        }

        public String getQueryHandleString() {
            return ctx.getQueryHandleString();
        }

        public QueryHandle getQueryHandle() {
            return ctx.getQueryHandle();
        }
    }

    /**
     * The Class QuerySubmitter.
     */
    private class QuerySubmitter implements Runnable {

        /**
         * The paused for test.
         */
        private boolean pausedForTest = false;

        private final EstimatedQueryCollection waitingQueries;

        private final QueryLaunchingConstraintsChecker constraintsChecker;

        public QuerySubmitter(@NonNull final EstimatedQueryCollection waitingQueries,
                @NonNull final QueryLaunchingConstraintsChecker constraintsChecker) {

            this.waitingQueries = waitingQueries;
            this.constraintsChecker = constraintsChecker;
        }

        /*
         * (non-Javadoc)
         *
         * @see java.lang.Runnable#run()
         */
        @Override
        public void run() {
            log.info("Starting QuerySubmitter thread");
            while (!stopped && !querySubmitter.isInterrupted()) {
                try {
                    if (pausedForTest) {
                        Thread.sleep(100);
                        continue;
                    }
                    final QueryContext query = queuedQueries.take();
                    synchronized (query) {
                        /* Setting log segregation id */
                        logSegregationContext.setLogSegragationAndQueryId(query.getQueryHandleString());

                        if (!query.queued()) {
                            log.info("Probably the query got cancelled. Skipping it. Query Status:{}",
                                    query.getStatus());
                            continue;
                        }

                        log.info("Processing query:{}", query.getUserQuery());
                        /* Check javadoc of QueryExecutionServiceImpl#removalFromLaunchedQueriesLock for reason for existence
                         of this lock. */
                        log.debug("Acquiring lock in QuerySubmitter");
                        removalFromLaunchedQueriesLock.lock();
                        try {
                            if (this.constraintsChecker.canLaunch(query, launchedQueries)) {

                                /* Query is not going to be added to waiting queries. No need to keep the lock.
                                 First release lock, then launch query */
                                removalFromLaunchedQueriesLock.unlock();
                                // add it to launched queries data structure immediately sothat other constraint checks can start seeing
                                // this query
                                query.setLaunching(true);
                                launchedQueries.add(query);
                                Future launcherFuture = queryLauncherPool.submit(new QueryLauncher(query));
                                query.setQueryLauncher(launcherFuture);
                            } else {
                                /* Query is going to be added to waiting queries. Keep holding the lock to avoid any removal from
                                launched queries. First add to waiting queries, then release lock */
                                addToWaitingQueries(query);
                                removalFromLaunchedQueriesLock.unlock();
                                if (query.getRetryPolicy() != null) {
                                    waitingQueriesSelectionSvc.schedule(new Runnable() {
                                        @Override
                                        public void run() {
                                            if (waitingQueries.remove(query)) {
                                                queuedQueries.add(query);
                                            }
                                        }
                                    }, query.getRetryPolicy().getOperationNextTime(query)
                                            - System.currentTimeMillis(), TimeUnit.MILLISECONDS);
                                }
                            }
                        } finally {
                            if (removalFromLaunchedQueriesLock.isHeldByCurrentThread()) {
                                removalFromLaunchedQueriesLock.unlock();
                            }
                        }
                    }
                } catch (InterruptedException e) {
                    log.info("Query Submitter has been interrupted, exiting");
                    return;
                } catch (Exception e) {
                    incrCounter(QUERY_SUBMITTER_COUNTER);
                    log.error("Error in query submitter", e);
                }
            }
            log.info("QuerySubmitter exited");
        }

        private void addToWaitingQueries(final QueryContext query) throws LensException {

            checkEstimatedQueriesState(query);
            this.waitingQueries.add(query);
            log.info("Added to waiting queries. QueryId:{}", query.getQueryHandleString());
        }
    }

    private class QueryLauncher implements Runnable {
        QueryContext query;

        QueryLauncher(QueryContext query) {
            this.query = query;
            log.info("Query launcher created for query {} on driver {}", query.getQueryHandle(),
                    query.getSelectedDriver().getFullyQualifiedName());
            query.setLaunchTime(System.currentTimeMillis());
        }

        @Override
        public void run() {
            synchronized (query) {
                try {
                    logSegregationContext.setLogSegragationAndQueryId(query.getQueryHandleString());
                    log.info("Starting to launch query {} on driver {}", query.getQueryHandle(),
                            query.getSelectedDriver().getFullyQualifiedName());
                    // acquire session before launching query.
                    acquire(query.getLensSessionIdentifier());
                    if (!query.getStatus().cancelled()) {
                        launchQuery(query);
                    }
                } catch (Exception e) {
                    if (!query.getStatus().cancelled()) {
                        log.error("Error launching query: {}", query.getQueryHandle(), e);
                        incrCounter(QUERY_SUBMITTER_COUNTER);
                        try {
                            setFailedStatus(query, "Launching query failed", e);
                        } catch (LensException e1) {
                            log.error("Error in setting failed status", e1);
                        }
                    }
                } finally {
                    query.setLaunching(false);
                    try {
                        release(query.getLensSessionIdentifier());
                    } catch (LensException e) {
                        log.error("Error releasing session", e);
                    }
                }
            }
        }

        private void launchQuery(final QueryContext query) throws LensException {
            checkEstimatedQueriesState(query);
            query.getSelectedDriver().getQueryHook().preLaunch(query);
            QueryStatus oldStatus = query.getStatus();
            // If driver supports async updates.
            if (query.getSelectedDriver().getStatusUpdateMethod() == StatusUpdateMethod.PUSH) {
                query.registerStatusUpdateListener(asyncStatusUpdater);
            }
            QueryStatus newStatus = new QueryStatus(query.getStatus().getProgress(), null,
                    QueryStatus.Status.LAUNCHED, "Query is launched on driver", false, null, null, null);
            query.validateTransition(newStatus);
            // Check if we need to pass session's effective resources to selected driver
            addSessionResourcesToDriver(query);
            query.getSelectedDriver().executeAsync(query);
            query.setStatusSkippingTransitionTest(newStatus);
            query.clearTransientStateAfterLaunch();
            log.info("Added to launched queries. QueryId:{}", query.getQueryHandleString());
            fireStatusChangeEvent(query, newStatus, oldStatus);
        }
    }

    private void checkEstimatedQueriesState(final QueryContext query) throws LensException {
        if (query.getSelectedDriver() == null || query.getSelectedDriverQueryCost() == null) {
            throw new LensException("selected driver: " + query.getSelectedDriver()
                    + " OR selected driver query cost: " + query.getSelectedDriverQueryCost()
                    + " is null. Query doesn't appear to be an estimated query.");
        }
    }

    /**
     * Pause query submitter.
     * note : used in tests only
     */
    public void pauseQuerySubmitter(boolean pause) {
        querySubmitterRunnable.pausedForTest = pause;
    }

    private class AsyncStatusUpdater implements QueryDriverStatusUpdateListener {

        @Override
        public void onDriverStatusUpdated(QueryHandle handle, DriverQueryStatus status) {
            try {
                updateStatus(handle, false);
            } catch (LensException e) {
                log.error("Unable to update status from driver status for query {}", handle, e);
            }
        }
    }

    /**
     * The Class StatusPoller.
     */
    private class StatusPoller implements Runnable {

        /**
         * The poll interval.
         */
        long pollInterval = 1000;

        /*
         * (non-Javadoc)
         *
         * @see java.lang.Runnable#run()
         */
        @Override
        public void run() {
            log.info("Starting Status poller thread");
            while (!stopped && !statusPoller.isInterrupted()) {
                try {
                    Set<QueryContext> launched = launchedQueries.getQueries();

                    for (QueryContext ctx : launched) {
                        if (stopped || statusPoller.isInterrupted()) {
                            return;
                        }
                        if (ctx.isLaunching()) {
                            continue;
                        }
                        logSegregationContext.setLogSegragationAndQueryId(ctx.getQueryHandleString());
                        log.debug("Polling status for {}", ctx.getQueryHandle());
                        try {
                            // session is not required to update status of the query
                            // don't need to wrap this with acquire/release
                            updateStatus(ctx.getQueryHandle());
                        } catch (LensException e) {
                            log.error("Error updating status ", e);
                        }
                    }
                    Thread.sleep(pollInterval);
                } catch (InterruptedException e) {
                    log.info("Status poller has been interrupted, exiting");
                    return;
                } catch (Exception e) {
                    incrCounter(STATUS_UPDATE_COUNTER);
                    log.error("Error in status poller", e);
                }
            }
            log.info("StatusPoller exited");
        }
    }

    private boolean handleRetries(QueryContext ctx) throws LensException {
        // TODO: handle retries for post-processing, e.g. result formatting failure doesn't need query rerun
        if (ctx.getStatus().failing()) {
            if (removeFromLaunchedQueries(ctx)) {
                processWaitingQueriesAsync(ctx);
            }
            if (ctx.getDriverStatus().failed() && !getDriverRetryPolicy(ctx).hasExhaustedRetries(ctx)) {
                log.info("query {} will be retried on the same driver {}", ctx.getQueryHandle(),
                        ctx.getSelectedDriver().getFullyQualifiedName());
                ctx.extractFailedAttempt();
                ctx.setStatus(QueryStatus.getQueuedStatus());
                ctx.getSelectedDriver().closeQuery(ctx.getQueryHandle());
                return queuedQueries.add(ctx);
            } else if (!getServerRetryPolicy(ctx).hasExhaustedRetries(ctx)) {
                LensDriver selectedDriver = ctx.getSelectedDriver();
                ctx.getDriverContext().blacklist(selectedDriver);
                try (SessionContext ignored = new SessionContext(
                        getSessionHandle(ctx.getLensSessionIdentifier()))) {
                    rewriteAndSelect(ctx);
                } catch (LensException e) {
                    log.error(
                            "driver {} gave up on query {} and it will not be retried on any other driver since rewrite failed",
                            selectedDriver.getFullyQualifiedName(), e);
                    ctx.setStatus(new QueryStatus(1.0f, null, FAILED, ctx.getStatus().getStatusMessage(), false,
                            null, ctx.getStatus().getErrorMessage(), ctx.getStatus().getLensErrorTO()));
                    return false;
                }
                log.info("driver {} gave up on query {} and it will be retried on {}",
                        selectedDriver.getFullyQualifiedName(), ctx.getQueryHandle(),
                        ctx.getSelectedDriver().getFullyQualifiedName());
                ctx.extractFailedAttempt(selectedDriver);
                ctx.setStatus(QueryStatus.getQueuedStatus());
                selectedDriver.closeQuery(ctx.getQueryHandle());
                return queuedQueries.add(ctx);
            }
            ctx.setStatus(new QueryStatus(1.0f, null, FAILED, ctx.getStatus().getStatusMessage(), false, null,
                    ctx.getStatus().getErrorMessage(), ctx.getStatus().getLensErrorTO()));
        }
        return false;
    }

    /**
     * Sets the failed status.
     *
     * @param ctx       the ctx
     * @param statusMsg the status msg
     * @param e    the LensException
     * @throws LensException the lens exception
     */
    void setFailedStatus(QueryContext ctx, String statusMsg, Exception e) throws LensException {
        QueryStatus before = ctx.getStatus();
        ctx.setStatus(new QueryStatus(0.0f, null, FAILING, statusMsg, false, null, LensUtil.getCauseMessage(e),
                e instanceof LensException ? ((LensException) e).buildLensErrorTO(this.errorCollection) : null));
        handleRetries(ctx);
        if (ctx.finished()) {
            updateFinishedQuery(ctx, before);
        }
        fireStatusChangeEvent(ctx, ctx.getStatus(), before);
    }

    private BackOffRetryHandler<QueryContext> getServerRetryPolicy(QueryContext ctx) {
        if (ctx.getServerRetryPolicy() == null) {
            // allow new driver to retry
            ctx.setDriverRetryPolicy(null);
            ctx.setServerRetryPolicy(queryRetryPolicyDecider.decidePolicy(ctx.getStatus().getErrorMessage()));
        }
        return ctx.getServerRetryPolicy();
    }

    private BackOffRetryHandler<QueryContext> getDriverRetryPolicy(QueryContext ctx) {
        if (ctx.getDriverRetryPolicy() == null) {
            ctx.setDriverRetryPolicy(ctx.getSelectedDriver().getRetryPolicyDecider()
                    .decidePolicy(ctx.getDriverStatus().getErrorMessage()));
        }
        return ctx.getDriverRetryPolicy();
    }

    /**
     * Sets the cancelled status.
     *
     * @param ctx       the ctx
     * @param statusMsg the status msg
     * @throws LensException the lens exception
     */
    private void setCancelledStatus(QueryContext ctx, String statusMsg) throws LensException {
        QueryStatus before = ctx.getStatus();
        ctx.setStatus(new QueryStatus(0.0f, null, CANCELED, statusMsg, false, null, null, null));
        updateFinishedQuery(ctx, before);
        fireStatusChangeEvent(ctx, ctx.getStatus(), before);
    }

    /**
     * Update finished query.
     *
     * @param ctx    the ctx
     * @param before the before
     */
    private void updateFinishedQuery(QueryContext ctx, QueryStatus before) {
        // before would be null in case of server restart
        if (before != null) {
            /* Seems like query is cancelled, remove it from both queuedQueries and waitingQueries because we don't know
             * where it is right now. It might happen that when we remove it from queued, it was in waiting OR
             * when we removed it from waiting, it was in queued. We might just miss removing it from everywhere due to this
             * hide and seek. Then QuerySubmitter thread will come to rescue, as it always checks that a query should be in
             * queued state before processing it after deque. If it is in cancelled state, then it will skip it. */
            queuedQueries.remove(ctx);
            waitingQueries.remove(ctx);
            // Remove it from launched queries as well - the query got cancelled or failed to launch or failed to execute
            if (removeFromLaunchedQueries(ctx)) {
                processWaitingQueriesAsync(ctx);
            }
        }
        // Remove from active queries
        if (SESSION_MAP.containsKey(ctx.getLensSessionIdentifier())) {
            getSession(SESSION_MAP.get(ctx.getLensSessionIdentifier()))
                    .removeFromActiveQueries(ctx.getQueryHandle());
        }
        finishedQueries.add(new FinishedQuery(ctx));
        ctx.clearTransientStateAfterLaunch();
    }

    void setSuccessState(QueryContext ctx) throws LensException {
        QueryStatus before = ctx.getStatus();
        ctx.setStatus(new QueryStatus(1.0f, null, SUCCESSFUL, "Query is successful!",
                ctx.isResultAvailableInDriver(), null, null, null));
        updateFinishedQuery(ctx, before);
        fireStatusChangeEvent(ctx, ctx.getStatus(), before);
    }

    /**
     * Update status.
     *
     * @param handle the handle
     * @throws LensException the lens exception
     */
    private void updateStatus(final QueryHandle handle) throws LensException {
        updateStatus(handle, true);
    }

    private void updateStatus(final QueryHandle handle, boolean updateDriverStatus) throws LensException {
        QueryContext ctx = allQueries.get(handle);
        if (ctx != null) {
            logSegregationContext.setLogSegragationAndQueryId(ctx.getLogHandle());
            log.info("Updating status for {}", ctx.getQueryHandle());
            synchronized (ctx) {
                QueryStatus before = ctx.getStatus();
                if (!ctx.queued() && !ctx.finished()) {
                    if (updateDriverStatus) {
                        try {
                            ctx.updateDriverStatus(statusUpdateRetryHandler);
                        } catch (LensException exc) {
                            // Status update from driver failed
                            setFailedStatus(ctx, "Status update failed", exc);
                            log.error("Status update failed for {}", handle, exc);
                            return;
                        }
                    }
                    ctx.setStatus(ctx.getDriverStatus().toQueryStatus());
                    // query is successfully executed by driver and
                    // if query result need not be persisted or there is no result available in driver, move the query to
                    // succeeded state immediately, otherwise result formatter will format the result and move it to succeeded
                    if (ctx.getStatus().getStatus().equals(EXECUTED)
                            && (!ctx.isPersistent() || !ctx.isResultAvailableInDriver())) {
                        setSuccessState(ctx);
                    } else {
                        if (ctx.getStatus().failing()) {
                            handleRetries(ctx);
                        }
                        if (ctx.getStatus().finished()) {
                            updateFinishedQuery(ctx, before);
                        }
                        fireStatusChangeEvent(ctx, ctx.getStatus(), before);
                    }
                }
                if (ctx.queued()) {
                    Integer queryIndex = waitingQueries.getQueryIndex(ctx);
                    // Query index could be null when the query status is queued but
                    // query is present in priorityblocking queue for processing
                    if (queryIndex != null) {
                        ctx.getStatus().setQueueNumber(queryIndex);
                    }
                }
            }
        }
    }

    /**
     * New status change event.
     *
     * @param ctx       the ctx
     * @param prevState the prev state
     * @param currState the curr state
     * @return the status change
     */
    private static StatusChange newStatusChangeEvent(QueryContext ctx, QueryStatus.Status prevState,
            QueryStatus.Status currState) {
        QueryHandle query = ctx.getQueryHandle();
        switch (currState) {
        case CANCELED:
            return new QueryCancelled(ctx, prevState, currState, null);
        case CLOSED:
            return new QueryClosed(ctx, prevState, currState, null);
        case FAILED:
            StringBuilder msgBuilder = new StringBuilder();
            msgBuilder.append(ctx.getStatus().getStatusMessage());
            if (!StringUtils.isBlank(ctx.getStatus().getErrorMessage())) {
                msgBuilder.append("\n Reason:\n");
                msgBuilder.append(ctx.getStatus().getErrorMessage());
            }
            return new QueryFailed(ctx, prevState, currState, msgBuilder.toString());
        case LAUNCHED:
            return new QueryLaunched(ctx.getLaunchTime(), prevState, currState, query);
        case QUEUED:
            if (ctx.getFailedAttempts().size() > 0) {
                return new QueryQueuedForRetry(ctx.getSubmissionTime(), prevState, currState, query,
                        ctx.getSubmittedUser());
            } else {
                return new QueryQueued(ctx.getSubmissionTime(), prevState, currState, query,
                        ctx.getSubmittedUser());
            }
        case RUNNING:
            return new QueryRunning(System.currentTimeMillis() - ctx.getDriverStatus().getDriverStartTime(),
                    prevState, currState, query);
        case EXECUTED:
            return new QueryExecuted(ctx.getDriverStatus().getDriverFinishTime(), prevState, currState, query);
        case SUCCESSFUL:
            return new QuerySuccess(ctx, prevState, currState);
        default:
            log.warn("Query {} transitioned to {} state from {} state", query, currState, prevState);
            return null;
        }
    }

    /**
     * If query status has changed, fire a specific StatusChange event.
     *
     * @param ctx     the ctx
     * @param current the current
     * @param before  the before
     */
    private void fireStatusChangeEvent(QueryContext ctx, QueryStatus current, QueryStatus before) {
        if (ctx == null || current == null) {
            return;
        }

        QueryStatus.Status prevState = before.getStatus();
        QueryStatus.Status currentStatus = current.getStatus();
        if (currentStatus.equals(prevState)) {
            // No need to fire event since the state hasn't changed
            return;
        }

        StatusChange event = newStatusChangeEvent(ctx, prevState, currentStatus);
        if (event != null) {
            try {
                getEventService().notifyEvent(event);
            } catch (LensException e) {
                log.warn("LensEventService encountered error while handling event: {}", event.getEventId(), e);
            }
        }
    }

    /**
     * The Class QueryPurger.
     */
    private class QueryPurger implements Runnable {

        /*
         * (non-Javadoc)
         *
         * @see java.lang.Runnable#run()
         */
        @Override
        public void run() {
            log.info("Starting Query purger thread");
            while (!stopped && !queryPurger.isInterrupted()) {
                try {
                    Iterator<FinishedQuery> iter = finishedQueries.iterator();
                    FinishedQuery finished;
                    while (iter.hasNext()) {
                        finished = iter.next();
                        if (finished.canBePurged()) {
                            try {
                                persistQuery(finished);
                                iter.remove();
                            } catch (Exception e) {
                                log.warn("Exception while purging query {}", finished.getQueryHandle(), e);
                                continue;
                            } finally {
                                if (!finished.getCtx().isQueryClosedOnDriver()) {
                                    try {
                                        if (finished.getCtx().getSelectedDriver() != null) {
                                            finished.getCtx().getSelectedDriver()
                                                    .closeQuery(finished.getQueryHandle());
                                        }
                                        finished.getCtx().setQueryClosedOnDriver(true);
                                    } catch (Exception e) {
                                        log.warn("Exception while closing query with selected driver.", e);
                                    }
                                    processWaitingQueriesAsync(finished.ctx);
                                }
                            }
                            synchronized (finished.ctx) {
                                finished.ctx.setFinishedQueryPersisted(true);
                                log.info("Purging: {}", finished.getQueryHandle());
                                allQueries.remove(finished.getQueryHandle());
                                resultSets.remove(finished.getQueryHandle());
                            }
                            fireStatusChangeEvent(finished.getCtx(),
                                    new QueryStatus(1f, null, CLOSED, "Query purged", false, null, null, null),
                                    finished.getCtx().getStatus());
                            log.info("Query purged: {}", finished.getQueryHandle());
                        }
                    }
                    Thread.sleep(purgeInterval);
                } catch (InterruptedException e) {
                    log.error("purger interrupted", e);
                } catch (Throwable e) {
                    log.error("Purger giving error", e);
                    incrCounter(QUERY_PURGER_COUNTER);
                }
            }
            log.info("QueryPurger exited");
        }

        private void persistQuery(FinishedQuery finished) throws SQLException {
            FinishedLensQuery finishedQuery = new FinishedLensQuery(finished.getCtx());
            if (finished.ctx.getStatus().getStatus() == SUCCESSFUL) {
                if (finished.ctx.getStatus().isResultSetAvailable()) {
                    try {
                        LensResultSet set = finished.getResultset();
                        if (set != null && PersistentResultSet.class.isAssignableFrom(set.getClass())) {
                            LensResultSetMetadata metadata = set.getMetadata();
                            String outputPath = set.getOutputPath();
                            Long fileSize = ((PersistentResultSet) set).getFileSize();
                            Integer rows = set.size();
                            finishedQuery.setResult(outputPath);
                            finishedQuery.setMetadata(metadata.toJson());
                            finishedQuery.setRows(rows);
                            finishedQuery.setFileSize(fileSize);
                        }
                    } catch (Exception e) {
                        log.error(
                                "Couldn't obtain result set info for the query: {}. Going ahead with perstsiting the query",
                                finished.getQueryHandle(), e);
                    }
                }
            }
            lensServerDao.insertFinishedQuery(finishedQuery);
            log.info("Saved query {} to DB", finishedQuery.getHandle());
        }
    }

    /**
     * The Class PreparedQueryPurger.
     */
    private class PreparedQueryPurger implements Runnable {

        /*
         * (non-Javadoc)
         *
         * @see java.lang.Runnable#run()
         */
        @Override
        public void run() {
            log.info("Starting Prepared Query purger thread");
            while (!stopped && !prepareQueryPurger.isInterrupted()) {
                try {
                    PreparedQueryContext prepared = preparedQueryQueue.take();
                    logSegregationContext.setLogSegragationAndQueryId(prepared.getQueryHandleString());
                    destroyPreparedQuery(prepared);
                    log.info("Purged prepared query: {}", prepared.getPrepareHandle());
                } catch (LensException e) {
                    incrCounter(PREPARED_QUERY_PURGER_COUNTER);
                    log.error("Error closing prepared query ", e);
                } catch (InterruptedException e) {
                    log.info("PreparedQueryPurger has been interrupted, exiting");
                    return;
                } catch (Exception e) {
                    incrCounter(PREPARED_QUERY_PURGER_COUNTER);
                    log.error("Error in prepared query purger", e);
                }
            }
            log.info("PreparedQueryPurger exited");
        }
    }

    /*
     * (non-Javadoc)
     *
     * @see org.apache.hive.service.CompositeService#init(org.apache.hadoop.hive.conf.HiveConf)
     */
    public synchronized void init(HiveConf hiveConf) {
        super.init(hiveConf);
        this.conf = hiveConf;
        this.isDuplicateQueryAllowed = conf.getBoolean(LensConfConstants.SERVER_DUPLICATE_QUERY_ALLOWED,
                LensConfConstants.DEFAULT_SERVER_DUPLICATE_QUERY_ALLOWED);
        try {
            loadQueryComparator();
        } catch (LensException e) {
            log.error("Error while loading query comparator class", e);
            throw new IllegalStateException("Could not load query comparator class", e);
        }

        this.launchedQueries = new ThreadSafeEstimatedQueryCollection(
                new DefaultEstimatedQueryCollection(new DefaultQueryCollection()));
        this.queuedQueries = new FairPriorityBlockingQueue<QueryContext>(queryComparator);

        this.waitingQueries = new ThreadSafeEstimatedQueryCollection(new DefaultEstimatedQueryCollection(
                new DefaultQueryCollection(new TreeSet<QueryContext>(queryComparator))));

        ImmutableSet<QueryLaunchingConstraint> queryConstraints = getImplementations(
                QUERY_LAUNCHING_CONSTRAINT_FACTORIES_KEY, hiveConf);

        this.queryConstraintsChecker = new DefaultQueryLaunchingConstraintsChecker(queryConstraints);

        this.querySubmitterRunnable = new QuerySubmitter(this.waitingQueries, this.queryConstraintsChecker);
        this.querySubmitter = new Thread(querySubmitterRunnable, "QuerySubmitter");

        ImmutableSet<WaitingQueriesSelectionPolicy> selectionPolicies = getImplementations(
                WAITING_QUERIES_SELECTION_POLICY_FACTORIES_KEY, hiveConf);

        this.waitingQueriesSelector = new UnioningWaitingQueriesSelector(selectionPolicies);

        try {
            this.userQueryToCubeQueryRewriter = new UserQueryToCubeQueryRewriter(conf);
        } catch (LensException e) {
            throw new IllegalStateException("Could not load phase 1 rewriters");
        }
        try {
            loadQueryRetryPolicyDecider(conf);
        } catch (LensException e) {
            throw new IllegalStateException("Could not load retry policy", e);
        }
        try {
            initializeQueryAcceptors();
        } catch (LensException e) {
            throw new IllegalStateException("Could not load acceptors");
        }
        initializeListeners();
        try {
            // Load all configured Drivers
            loadDrivers();
            // load driver selector
            loadDriverSelector();
        } catch (LensException e) {
            log.error("Error while loading drivers", e);
            throw new IllegalStateException("Could not load drivers", e);
        }
        purgeInterval = conf.getInt(PURGE_INTERVAL, DEFAULT_PURGE_INTERVAL);
        initalizeFinishedQueryStore(conf);

        inMemoryResultsetTTLMillis = conf.getInt(LensConfConstants.INMEMORY_RESULT_SET_TTL_SECS,
                LensConfConstants.DEFAULT_INMEMORY_RESULT_SET_TTL_SECS) * 1000;

        int statusUpdateRetries = conf.getInt(LensConfConstants.STATUS_UPDATE_EXPONENTIAL_RETRIES,
                LensConfConstants.DEFAULT_STATUS_UPDATE_EXPONENTIAL_RETRIES);
        //  Maximum delay a status update can wait for next update, in case of transient failures.
        long statusUpdateRetryMaxDelay = conf.getLong(LensConfConstants.MAXIMUM_STATUS_UPDATE_DELAY,
                LensConfConstants.DEFAULT_MAXIMUM_STATUS_UPDATE_DELAY) * 1000;
        // The wait time for next status update which can grow exponentially, in case of transient failures.
        long statusUpdateExponentialWaiFactor = conf.getLong(
                LensConfConstants.STATUS_UPDATE_EXPONENTIAL_WAIT_FACTOR,
                LensConfConstants.DEFAULT_STATUS_UPDATE_EXPONENTIAL_WAIT_FACTOR);
        statusUpdateRetryHandler = OperationRetryHandlerFactory.createExponentialBackOffHandler(statusUpdateRetries,
                statusUpdateRetryMaxDelay, statusUpdateExponentialWaiFactor);
        log.info("Query execution service initialized");
    }

    private void loadQueryRetryPolicyDecider(Configuration conf) throws LensException {
        this.queryRetryPolicyDecider = ChainedRetryPolicyDecider.from(conf, QUERY_RETRY_POLICY_CLASSES);
    }

    /**
     * Initalize finished query store.
     *
     * @param conf the conf
     */
    private void initalizeFinishedQueryStore(Configuration conf) {
        this.lensServerDao = new LensServerDAO();
        this.lensServerDao.init(conf);
        try {
            this.lensServerDao.createFinishedQueriesTable();
            this.lensServerDao.createFailedAttemptsTable();
        } catch (Exception e) {
            log.warn("Unable to create finished query tables, query purger will not purge queries", e);
        }
    }

    private void awaitTermination(ExecutorService service) {
        try {
            service.awaitTermination(1, TimeUnit.MINUTES);
        } catch (InterruptedException e) {
            log.info("Couldn't finish executor service within 1 minute: {}", service);
        }
    }

    private void awaitTermination(QueryResultPurger service) {
        if (service == null) {
            return;
        }
        try {
            service.awaitTermination(1, TimeUnit.MINUTES);
        } catch (InterruptedException e) {
            log.info("Couldn't finish query result purger within 1 minute: {}", service);
        }
    }

    /*
     * (non-Javadoc)
     *
     * @see org.apache.lens.server.LensService#prepareStopping()
     */
    public void prepareStopping() {
        super.prepareStopping();
        Thread[] threadsToStop = new Thread[] { querySubmitter, statusPoller, queryPurger, prepareQueryPurger };
        // Nudge the threads to stop
        for (Thread th : threadsToStop) {
            th.interrupt();
        }

        // Nudge executor pools to stop

        // Hard shutdown, since it doesn't matter whether waiting queries were selected, all will be
        // selected in the next restart
        waitingQueriesSelectionSvc.shutdownNow();
        // Soft shutdown, Wait for current estimate tasks
        estimatePool.shutdown();
        // shutdown launcher pool
        queryLauncherPool.shutdown();
        // Soft shutdown for result purger too. Purging shouldn't take much time.
        if (null != queryResultPurger) {
            queryResultPurger.shutdown();
        }
        // shutdown query expirer
        queryExpirer.shutdownNow();
        // Soft shutdown right now, will await termination in this method itself, since cancellation pool
        // should be terminated before query state gets persisted.
        queryCancellationPool.shutdown();

        // Join the threads.
        for (Thread th : threadsToStop) {
            try {
                log.debug("Waiting for {}", th.getName());
                th.join();
            } catch (InterruptedException e) {
                log.error("Error waiting for thread: {}", th.getName(), e);
            }
        }
        // Needs to be done before queries' states are persisted, hence doing here. Await of other
        // executor services can be done after persistence, hence they are done in #stop
        awaitTermination(queryLauncherPool);
        awaitTermination(queryCancellationPool);
    }

    /*
     * (non-Javadoc)
     *
     * @see org.apache.hive.service.CompositeService#stop()
     */
    public synchronized void stop() {
        super.stop();
        awaitTermination(waitingQueriesSelectionSvc);
        awaitTermination(estimatePool);
        awaitTermination(queryResultPurger);
        log.info("Query execution service stopped");
    }

    /*
     * (non-Javadoc)
     *
     * @see org.apache.hive.service.CompositeService#start()
     */
    public synchronized void start() {
        synchronized (allQueries) {
            // populate the query queues
            final List<QueryContext> allRestoredQueuedQueries = new LinkedList<QueryContext>();
            for (QueryContext ctx : allQueries.values()) {
                // recover query configurations from session
                try {
                    if (SESSION_MAP.containsKey(ctx.getLensSessionIdentifier())) {
                        // try setting configuration if the query session is still not closed
                        ctx.setConf(
                                getLensConf(getSessionHandle(ctx.getLensSessionIdentifier()), ctx.getLensConf()));
                        // Add queryHandle to active queries
                        getSession(SESSION_MAP.get(ctx.getLensSessionIdentifier()))
                                .addToActiveQueries(ctx.getQueryHandle());
                    } else {
                        ctx.setConf(getLensConf(ctx.getLensConf()));
                    }
                    for (LensDriver driver : drivers.values()) {
                        if (ctx.getDriverContext() != null) {
                            ctx.getDriverContext().setDriverConf(driver, ctx.getConf());
                        }
                    }
                } catch (LensException e) {
                    log.error("Could not set query conf ", e);
                }
                // Add queries to the queue.
                switch (ctx.getStatus().getStatus()) {
                case NEW:
                case QUEUED:
                    allRestoredQueuedQueries.add(ctx);
                    break;
                case LAUNCHED:
                case RUNNING:
                case EXECUTED:
                    try {
                        launchedQueries.add(ctx);
                    } catch (final Exception e) {
                        log.error("Query not restored:QueryContext:{}", ctx, e);
                    }
                    // If EXECUTED, try to nudge result formatting forward
                    if (ctx.getStatus().getStatus() == EXECUTED) {
                        try {
                            getEventService()
                                    .notifyEvent(newStatusChangeEvent(ctx, null, ctx.getStatus().getStatus()));
                        } catch (LensException e) {
                            log.error("Couldn't notify event for query executed for {}", ctx, e);
                        }
                    }
                    break;
                case SUCCESSFUL:
                case FAILED:
                case CANCELED:
                    updateFinishedQuery(ctx, null);
                    break;
                case CLOSED:
                    allQueries.remove(ctx.getQueryHandle());
                    if (SESSION_MAP.containsKey(ctx.getLensSessionIdentifier())) {
                        getSession(SESSION_MAP.get(ctx.getLensSessionIdentifier()))
                                .removeFromActiveQueries(ctx.getQueryHandle());
                    }
                    log.info("Removed closed query from all Queries:" + ctx.getQueryHandle());
                }
            }
            queuedQueries.addAll(allRestoredQueuedQueries);
            log.info("Recovered {} queries", allQueries.size());
        }
        super.start();

        startEstimatePool();
        startLauncherPool();
        startQueryCancellationPool();

        querySubmitter.start();
        statusPoller.start();
        queryPurger.start();
        prepareQueryPurger.start();
        startQueryExpirer();

        if (conf.getBoolean(RESULTSET_PURGE_ENABLED, DEFAULT_RESULTSET_PURGE_ENABLED)) {
            queryResultPurger = new QueryResultPurger();
            queryResultPurger.init(conf);
        } else {
            log.info("Query result purger is not enabled");
        }
    }

    private void startEstimatePool() {
        int minPoolSize = conf.getInt(ESTIMATE_POOL_MIN_THREADS, DEFAULT_ESTIMATE_POOL_MIN_THREADS);
        int maxPoolSize = conf.getInt(ESTIMATE_POOL_MAX_THREADS, DEFAULT_ESTIMATE_POOL_MAX_THREADS);
        int keepAlive = conf.getInt(ESTIMATE_POOL_KEEP_ALIVE_MILLIS, DEFAULT_ESTIMATE_POOL_KEEP_ALIVE_MILLIS);

        final ThreadFactory defaultFactory = Executors.defaultThreadFactory();
        final AtomicInteger thId = new AtomicInteger();
        // We are creating our own thread factory, just so that we can override thread name for easy debugging
        ThreadFactory threadFactory = new ThreadFactory() {
            @Override
            public Thread newThread(Runnable r) {
                Thread th = defaultFactory.newThread(r);
                th.setName("estimate-" + thId.incrementAndGet());
                return th;
            }
        };

        log.debug("starting estimate pool");

        ThreadPoolExecutor estimatePool = new ThreadPoolExecutor(minPoolSize, maxPoolSize, keepAlive,
                TimeUnit.MILLISECONDS, new SynchronousQueue<Runnable>(), threadFactory);
        estimatePool.allowCoreThreadTimeOut(false);
        estimatePool.prestartCoreThread();
        this.estimatePool = estimatePool;
    }

    private void startLauncherPool() {
        int minPoolSize = conf.getInt(LAUNCHER_POOL_MIN_THREADS, DEFAULT_LAUNCHER_POOL_MIN_THREADS);
        int maxPoolSize = conf.getInt(LAUNCHER_POOL_MAX_THREADS, DEFAULT_LAUNCHER_POOL_MAX_THREADS);
        int keepAlive = conf.getInt(LAUNCHER_POOL_KEEP_ALIVE_MILLIS, DEFAULT_LAUNCHER_POOL_KEEP_ALIVE_MILLIS);

        final ThreadFactory defaultFactory = Executors.defaultThreadFactory();
        final AtomicInteger thId = new AtomicInteger();
        // We are creating our own thread factory, just so that we can override thread name for easy debugging
        ThreadFactory threadFactory = new ThreadFactory() {
            @Override
            public Thread newThread(Runnable r) {
                Thread th = defaultFactory.newThread(r);
                th.setName("launcher-" + thId.incrementAndGet());
                return th;
            }
        };

        log.debug("starting query launcher pool");

        ThreadPoolExecutor launcherPool = new ThreadPoolExecutor(minPoolSize, maxPoolSize, keepAlive,
                TimeUnit.MILLISECONDS, new SynchronousQueue<Runnable>(), threadFactory);
        launcherPool.allowCoreThreadTimeOut(false);
        launcherPool.prestartCoreThread();
        this.queryLauncherPool = launcherPool;
    }

    private void startQueryCancellationPool() {
        ThreadFactory factory = new BasicThreadFactory.Builder().namingPattern("query-cancellation-pool-Thread-%d")
                .priority(Thread.NORM_PRIORITY).build();
        //Using fixed values for pool . corePoolSize = maximumPoolSize = 3  and keepAliveTime = 60 secs
        queryCancellationPool = new ThreadPoolExecutor(3, 3, 60, TimeUnit.SECONDS, new LinkedBlockingQueue(),
                factory);
    }

    private void startQueryExpirer() {
        ThreadFactory factory = new BasicThreadFactory.Builder().namingPattern("QueryExpirer-%d").daemon(true)
                .priority(Thread.NORM_PRIORITY).build();
        queryExpirer = Executors.newSingleThreadScheduledExecutor(factory);
        long expiryRunInterval = conf.getLong(QUERY_EXPIRY_INTERVAL_MILLIS, DEFAULT_QUERY_EXPIRY_INTERVAL_MILLIS);
        queryExpirer.scheduleWithFixedDelay(new Runnable() {
            @Override
            public void run() {
                try {
                    expireQueries();
                } catch (Exception e) {
                    incrCounter(QUERY_EXPIRY_FAILURE_COUNTER);
                    log.error("Unable to expire queries", e);
                }
            }
        }, expiryRunInterval, expiryRunInterval, TimeUnit.MILLISECONDS);
        log.info("Enabled periodic exipry of queries at {} millis interval", expiryRunInterval);
    }

    private void expireQueries() {
        for (QueryContext ctx : launchedQueries.getQueries()) {
            try {
                if (ctx.hasTimedout()) {
                    log.info("Query {} has timedout, thus cancelling it", ctx.getLogHandle());
                    queryCancellationPool.submit(new CancelQueryTask(ctx.getQueryHandle()));
                    incrCounter(TOTAL_QUERIES_EXPIRED);
                }
            } catch (Exception e) {
                incrCounter(QUERY_EXPIRY_FAILURE_COUNTER);
                log.error("Unable to expire queries", e);
            }
        }
    }

    @AllArgsConstructor
    private class CancelQueryTask implements Runnable {
        private QueryHandle handle;

        @Override
        public void run() {
            try {
                logSegregationContext.setLogSegragationAndQueryId(handle.getHandleIdString());
                cancelQuery(handle);
            } catch (Exception e) {
                log.error("Error while cancelling query {}", handle, e);
            }
        }
    }

    private static final String REWRITE_GAUGE = "CUBE_REWRITE";
    private static final String DRIVER_ESTIMATE_GAUGE = "DRIVER_ESTIMATE";
    private static final String DRIVER_SELECTOR_GAUGE = "DRIVER_SELECTION";
    private static final String PARALLEL_CALL_GAUGE = "PARALLEL_ESTIMATE";

    /**
     * Rewrite the query for each driver, and estimate query cost for the rewritten queries. Finally, select the driver
     * using driver selector.
     *
     * @param ctx query context
     * @throws LensException the lens exception
     */
    void rewriteAndSelect(final AbstractQueryContext ctx) throws LensException {
        logSegregationContext.setLogSegragationAndQueryId(ctx.getLogHandle());
        log.info("Query details : User query: {} Submitted user: {} Session: {}", ctx.getUserQuery(),
                ctx.getSubmittedUser(), ctx.getLensSessionIdentifier());
        MethodMetricsContext parallelCallGauge = MethodMetricsFactory.createMethodGauge(ctx.getConf(), false,
                PARALLEL_CALL_GAUGE);
        try {
            userQueryToCubeQueryRewriter.rewrite(ctx);
            // Initially we obtain individual runnables for rewrite and estimate calls
            // These are mapped against the driver, so that later it becomes easy to chain them
            // for each driver.
            Map<LensDriver, RewriteUtil.DriverRewriterRunnable> rewriteRunnables = RewriteUtil.rewriteQuery(ctx);
            Map<LensDriver, AbstractQueryContext.DriverEstimateRunnable> estimateRunnables = ctx
                    .getDriverEstimateRunnables();

            int numDrivers = ctx.getDriverContext().getDrivers().size();
            final CountDownLatch estimateCompletionLatch = new CountDownLatch(numDrivers);
            List<RewriteEstimateRunnable> runnables = new ArrayList<RewriteEstimateRunnable>(numDrivers);
            List<Future> estimateFutures = new ArrayList<Future>();

            boolean cancelAllDueToOverload = false;
            for (final LensDriver driver : ctx.getDriverContext().getDrivers()) {
                RewriteEstimateRunnable r = new RewriteEstimateRunnable(driver, rewriteRunnables.get(driver),
                        estimateRunnables.get(driver), ctx, estimateCompletionLatch);

                // Submit composite rewrite + estimate operation to background pool
                try {
                    estimateFutures.add(estimatePool.submit(r));
                    runnables.add(r);
                } catch (RejectedExecutionException e) {
                    // This means the server can't accept anymore estimate requests at this time.
                    // Cancel all other futures.
                    cancelAllDueToOverload = true;
                    log.warn("Rejected from submitting to the estimate pool for driver {} ", r.getDriver(), e);
                    break;
                }
            }
            if (cancelAllDueToOverload) {
                for (int i = 0; i < runnables.size(); i++) {
                    RewriteEstimateRunnable r = runnables.get(i);
                    estimateFutures.get(i).cancel(true);
                    log.info("Cancelling estimate tasks for driver due to incomplete driver {}", r.getDriver());
                }
                throw new LensException(SERVER_OVERLOADED.getLensErrorInfo());
            }
            // Wait for all rewrite and estimates to finish
            try {
                long estimateLatchTimeout = ctx.getConf().getLong(ESTIMATE_TIMEOUT_MILLIS,
                        DEFAULT_ESTIMATE_TIMEOUT_MILLIS);
                boolean completed = estimateCompletionLatch.await(estimateLatchTimeout, TimeUnit.MILLISECONDS);

                // log operations yet to complete and  check if we can proceed with at least one driver
                if (!completed) {
                    int inCompleteDrivers = 0;

                    for (int i = 0; i < runnables.size(); i++) {
                        RewriteEstimateRunnable r = runnables.get(i);
                        if (!r.isCompleted()) {
                            ++inCompleteDrivers;
                            // Cancel the corresponding task
                            estimateFutures.get(i).cancel(true);
                            log.warn("Timeout reached for estimate task for driver {}" + r.getDriver());
                        }
                    }

                    if (inCompleteDrivers == ctx.getDriverContext().getDrivers().size()) {
                        throw new LensException(
                                "None of the drivers could complete within timeout: " + estimateLatchTimeout);
                    }
                }
            } catch (InterruptedException exc) {
                throw new LensException("At least one of the estimate operation failed to complete in time", exc);
            }

            // Evaluate success of rewrite and estimate
            boolean succeededOnce = false;
            List<String> failureCauses = new ArrayList<>(numDrivers);
            List<LensException> causes = new ArrayList<>(numDrivers);

            for (RewriteEstimateRunnable r : runnables) {
                if (r.isSucceeded()) {
                    succeededOnce = true;
                } else {
                    failureCauses.add(r.getFailureCause());

                    if (r.getCause() != null) {
                        causes.add(r.getCause());
                    }
                }
            }

            // Throw exception if none of the rewrite+estimates are successful.
            if (!succeededOnce) {
                if (!causes.isEmpty()) {
                    final LensException firstCause = causes.get(0);
                    for (LensException cause : causes) {
                        if (!cause.equals(firstCause)) {
                            throw new LensMultiCauseException(ImmutableList.copyOf(causes));
                        }
                    }
                    throw firstCause;
                } else {
                    throw new LensException(StringUtils.join(failureCauses, '\n'));
                }
            }

            MethodMetricsContext selectGauge = MethodMetricsFactory.createMethodGauge(ctx.getConf(), false,
                    DRIVER_SELECTOR_GAUGE);
            // 2. select driver to run the query
            LensDriver driver = driverSelector.select(ctx, conf);
            ctx.setSelectedDriver(driver);
            QueryCost selectedDriverQueryCost = ctx.getDriverContext().getDriverQueryCost(driver);
            ctx.setSelectedDriverQueryCost(selectedDriverQueryCost);
            Priority priority = driver.decidePriority(ctx);
            ctx.setPriority(priority == null ? Priority.NORMAL : priority);
            driver.getQueryHook().postDriverSelection(ctx);
            selectGauge.markSuccess();
        } finally {
            parallelCallGauge.markSuccess();
        }
    }

    /**
     * Chains driver specific rewrite and estimate of the query in a single runnable, which can be processed in a
     * background thread
     */
    public class RewriteEstimateRunnable implements Runnable {
        @Getter
        private final LensDriver driver;
        private final RewriteUtil.DriverRewriterRunnable rewriterRunnable;
        private final AbstractQueryContext.DriverEstimateRunnable estimateRunnable;
        private final AbstractQueryContext ctx;
        private final CountDownLatch estimateCompletionLatch;

        @Getter
        private boolean succeeded;
        @Getter
        private String failureCause = null;

        @Getter
        private LensException cause;

        @Getter
        private volatile boolean completed;

        public RewriteEstimateRunnable(LensDriver driver, RewriteUtil.DriverRewriterRunnable rewriterRunnable,
                AbstractQueryContext.DriverEstimateRunnable estimateRunnable, AbstractQueryContext ctx,
                CountDownLatch estimateCompletionLatch) {
            this.driver = driver;
            this.rewriterRunnable = rewriterRunnable;
            this.estimateRunnable = estimateRunnable;
            this.ctx = ctx;
            this.estimateCompletionLatch = estimateCompletionLatch;
        }

        @Override
        public void run() {
            try {
                // With following set - explain estimate calls are setting queryLogId as requestid in logSegregationContext
                logSegregationContext.setLogSegragationAndQueryId(ctx.getLogHandle());
                acquire(ctx.getLensSessionIdentifier());
                MethodMetricsContext rewriteGauge = MethodMetricsFactory
                        .createMethodGauge(ctx.getDriverConf(driver), true, REWRITE_GAUGE);
                log.info("Calling preRewrite hook for driver {}", driver.getFullyQualifiedName());
                driver.getQueryHook().preRewrite(ctx);
                // 1. Rewrite for driver
                rewriterRunnable.run();
                succeeded = rewriterRunnable.isSucceeded();
                if (!succeeded) {
                    failureCause = rewriterRunnable.getFailureCause();
                    cause = rewriterRunnable.getCause();
                } else {
                    log.info("Calling postRewrite hook for driver {}", driver.getFullyQualifiedName());
                    driver.getQueryHook().postRewrite(ctx);
                }

                rewriteGauge.markSuccess();

                // 2. Estimate for driver only if rewrite succeeded.
                if (succeeded) {
                    MethodMetricsContext estimateGauge = MethodMetricsFactory
                            .createMethodGauge(ctx.getDriverConf(driver), true, DRIVER_ESTIMATE_GAUGE);

                    log.info("Calling preEstimate hook for driver {}", driver.getFullyQualifiedName());
                    driver.getQueryHook().preEstimate(ctx);
                    estimateRunnable.run();
                    succeeded = estimateRunnable.isSucceeded();
                    if (!succeeded) {
                        failureCause = estimateRunnable.getFailureCause();
                        cause = estimateRunnable.getCause();
                        log.error("Estimate failed for driver {} cause: {}", driver, failureCause);
                    } else {
                        log.info("Calling postRewrite hook for driver {}", driver.getFullyQualifiedName());
                        driver.getQueryHook().postEstimate(ctx);
                    }

                    estimateGauge.markSuccess();
                } else {
                    log.error("Estimate skipped since rewrite failed for driver {} cause: {}", driver,
                            failureCause);
                }
            } catch (Throwable th) {
                log.error("Error computing estimate for driver {}", driver, th);
            } finally {
                completed = true;
                try {
                    release(ctx.getLensSessionIdentifier());
                } catch (LensException e) {
                    log.error("Could not release session: {}", ctx.getLensSessionIdentifier(), e);
                } finally {
                    estimateCompletionLatch.countDown();
                }
            }
        }
    }

    /**
     * Accept.
     *
     * @param query    the query
     * @param conf     the conf
     * @param submitOp the submit op
     * @throws LensException the lens exception
     */
    private void accept(String query, Configuration conf, SubmitOp submitOp) throws LensException {
        // run through all the query acceptors, and throw Exception if any of them
        // return false
        for (QueryAcceptor acceptor : queryAcceptors) {
            String cause = "";
            String rejectionCause = acceptor.accept(query, conf, submitOp);
            if (rejectionCause != null) {
                getEventService()
                        .notifyEvent(new QueryRejected(System.currentTimeMillis(), query, rejectionCause, null));
                throw new BadRequestException("Query not accepted because " + cause);
            }
        }
        getEventService().notifyEvent(new QueryAccepted(System.currentTimeMillis(), null, query, null));
    }

    /**
     * Gets the resultset from dao.
     *
     * @param queryHandle the query handle
     * @return the resultset from dao
     * @throws LensException the lens exception
     */
    private LensPersistentResult getResultsetFromDAO(QueryHandle queryHandle) throws LensException {
        FinishedLensQuery query = lensServerDao.getQuery(queryHandle.toString());
        if (query != null) {
            if (query.getResult() == null) {
                throw new NotFoundException("InMemory Query result purged " + queryHandle);
            }
            try {
                return new LensPersistentResult(query, conf);
            } catch (Exception e) {
                throw new LensException(e);
            }
        }
        throw new NotFoundException("Query not found: " + queryHandle);
    }

    /**
     * Gets the resultset.
     *
     * @param queryHandle the query handle
     * @return the resultset
     * @throws LensException the lens exception
     */
    LensResultSet getResultset(QueryHandle queryHandle) throws LensException {
        QueryContext ctx = allQueries.get(queryHandle);

        if (ctx == null) {
            return getResultsetFromDAO(queryHandle);
        } else {
            synchronized (ctx) {
                if (ctx.isFinishedQueryPersisted()) {
                    return getResultsetFromDAO(queryHandle);
                }
                if (ctx.successful()) { // Do not return any result set for queries that have not finished successfully.
                    LensResultSet resultSet = resultSets.get(queryHandle);
                    if (resultSet == null) {
                        if (ctx.isPersistent() && ctx.getQueryOutputFormatter() != null) {
                            resultSets.put(queryHandle, new LensPersistentResult(ctx, conf));
                        } else if (ctx.isResultAvailableInDriver() && !ctx.isQueryClosedOnDriver()) {
                            //InMemory result can not be returned for a closed query
                            resultSet = getDriverResultset(queryHandle);
                            resultSets.put(queryHandle, resultSet);
                        }
                    }
                }
            }

            LensResultSet result = resultSets.get(queryHandle);
            if (result == null) {
                throw new NotFoundException("Result set not available for query:" + queryHandle);
            }
            return result;
        }
    }

    /**
     * Gets the driver resultset.
     *
     * @param queryHandle the query handle
     * @return the driver resultset
     * @throws LensException the lens exception
     */
    LensResultSet getDriverResultset(QueryHandle queryHandle) throws LensException {
        return allQueries.get(queryHandle).getSelectedDriver().fetchResultSet(allQueries.get(queryHandle));
    }

    /*
     * (non-Javadoc)
     *
     * @see org.apache.lens.server.api.query.QueryExecutionService#prepare(org.apache.lens.api.LensSessionHandle,
     * java.lang.String, org.apache.lens.api.LensConf, java.lang.String)
     */
    @Override
    public QueryPrepareHandle prepare(LensSessionHandle sessionHandle, String query, LensConf lensConf,
            String queryName) throws LensException {
        PreparedQueryContext prepared = null;
        try {
            acquire(sessionHandle);
            prepared = prepareQuery(sessionHandle, query, lensConf, SubmitOp.PREPARE);
            prepared.setQueryName(queryName);
            prepared.getSelectedDriver().prepare(prepared);
            return prepared.getPrepareHandle();
        } catch (LensException e) {
            if (prepared != null) {
                destroyPreparedQuery(prepared);
            }
            throw e;
        } finally {
            release(sessionHandle);
        }
    }

    /**
     * Prepare query.
     *
     * @param sessionHandle the session handle
     * @param query         the query
     * @param lensConf      the lens conf
     * @param op            the op
     * @return the prepared query context
     * @throws LensException the lens exception
     */
    private PreparedQueryContext prepareQuery(LensSessionHandle sessionHandle, String query, LensConf lensConf,
            SubmitOp op) throws LensException {
        Configuration conf = getLensConf(sessionHandle, lensConf);
        accept(query, conf, op);
        PreparedQueryContext prepared = new PreparedQueryContext(query, getSession(sessionHandle).getLoggedInUser(),
                conf, lensConf, drivers.values());
        prepared.setLensSessionIdentifier(sessionHandle.getPublicId().toString());
        rewriteAndSelect(prepared);
        preparedQueries.put(prepared.getPrepareHandle(), prepared);
        preparedQueryQueue.add(prepared);
        incrCounter(PREPARED_QUERIES_COUNTER);
        return prepared;
    }

    /*
     * (non-Javadoc)
     *
     * @see
     * org.apache.lens.server.api.query.QueryExecutionService#explainAndPrepare(org.apache.lens.api.LensSessionHandle,
     * java.lang.String, org.apache.lens.api.LensConf, java.lang.String)
     */
    @Override
    public QueryPlan explainAndPrepare(LensSessionHandle sessionHandle, String query, LensConf lensConf,
            String queryName) throws LensException {
        PreparedQueryContext prepared = null;
        try {
            log.info("ExplainAndPrepare: session:{} query: {}", sessionHandle, query);
            acquire(sessionHandle);
            prepared = prepareQuery(sessionHandle, query, lensConf, SubmitOp.EXPLAIN_AND_PREPARE);
            prepared.setQueryName(queryName);
            addSessionResourcesToDriver(prepared);
            QueryPlan plan = prepared.getSelectedDriver().explainAndPrepare(prepared).toQueryPlan();
            plan.setPrepareHandle(prepared.getPrepareHandle());
            return plan;
        } catch (LensException e) {
            if (prepared != null) {
                destroyPreparedQuery(prepared);
            }
            throw e;
        } catch (UnsupportedEncodingException e) {
            if (prepared != null) {
                destroyPreparedQuery(prepared);
            }
            throw new LensException(e);
        } finally {
            release(sessionHandle);
        }
    }

    /*
     * (non-Javadoc)
     *
     * @see
     * org.apache.lens.server.api.query.QueryExecutionService#executePrepareAsync(org.apache.lens.api.LensSessionHandle,
     * org.apache.lens.api.query.QueryPrepareHandle, org.apache.lens.api.LensConf, java.lang.String)
     */
    @Override
    public QueryHandle executePrepareAsync(LensSessionHandle sessionHandle, QueryPrepareHandle prepareHandle,
            LensConf conf, String queryName) throws LensException {
        try {
            log.info("ExecutePrepareAsync: session:{} prepareHandle:{}", sessionHandle,
                    prepareHandle.getPrepareHandleId());
            acquire(sessionHandle);
            PreparedQueryContext pctx = getPreparedQueryContext(sessionHandle, prepareHandle);
            Configuration qconf = getLensConf(sessionHandle, conf);
            accept(pctx.getUserQuery(), qconf, SubmitOp.EXECUTE);
            QueryContext ctx = createContext(pctx, getSession(sessionHandle).getLoggedInUser(), conf, qconf, 0);
            if (StringUtils.isNotBlank(queryName)) {
                // Override previously set query name
                ctx.setQueryName(queryName);
            } else {
                ctx.setQueryName(pctx.getQueryName());
            }
            ctx.setLensSessionIdentifier(sessionHandle.getPublicId().toString());
            return submitQuery(ctx);
        } finally {
            release(sessionHandle);
        }
    }

    /*
     * (non-Javadoc)
     *
     * @see org.apache.lens.server.api.query.QueryExecutionService#executePrepare(org.apache.lens.api.LensSessionHandle,
     * org.apache.lens.api.query.QueryPrepareHandle, long, org.apache.lens.api.LensConf, java.lang.String)
     */
    @Override
    public QueryHandleWithResultSet executePrepare(LensSessionHandle sessionHandle,
            QueryPrepareHandle prepareHandle, long timeoutMillis, LensConf conf, String queryName)
            throws LensException {
        try {
            log.info("ExecutePrepare: session:{} prepareHandle: {} timeout:{}", sessionHandle,
                    prepareHandle.getPrepareHandleId(), timeoutMillis);
            acquire(sessionHandle);
            PreparedQueryContext pctx = getPreparedQueryContext(sessionHandle, prepareHandle);
            Configuration qconf = getLensConf(sessionHandle, conf);
            QueryContext ctx = createContext(pctx, getSession(sessionHandle).getLoggedInUser(), conf, qconf,
                    timeoutMillis);
            if (StringUtils.isNotBlank(queryName)) {
                // Override previously set query name
                ctx.setQueryName(queryName);
            } else {
                ctx.setQueryName(pctx.getQueryName());
            }
            return executeTimeoutInternal(sessionHandle, ctx, timeoutMillis, qconf);
        } finally {
            release(sessionHandle);
        }
    }

    /*
     * (non-Javadoc)
     *
     * @see org.apache.lens.server.api.query.QueryExecutionService#executeAsync(org.apache.lens.api.LensSessionHandle,
     * java.lang.String, org.apache.lens.api.LensConf, java.lang.String)
     */
    @Override
    public QueryHandle executeAsync(LensSessionHandle sessionHandle, String query, LensConf conf, String queryName)
            throws LensException {
        try {
            log.info("ExecuteAsync: session:{} query: {}", sessionHandle, query);
            acquire(sessionHandle);
            Configuration qconf = getLensConf(sessionHandle, conf);
            accept(query, qconf, SubmitOp.EXECUTE);
            if (!isDuplicateQueryAllowed) {
                QueryHandle previousHandle = checkForDuplicateQuery(query, sessionHandle, qconf, queryName);
                if (previousHandle != null) {
                    log.info("Query:{} Session:{} User:{} duplicate query found", query, sessionHandle,
                            getSession(sessionHandle).getLoggedInUser());
                    return previousHandle;
                }
            }
            QueryContext ctx = createContext(query, getSession(sessionHandle).getLoggedInUser(), conf, qconf, 0);
            // Should be set only once
            ctx.setQueryConfHash(UtilityMethods.generateHashOfWritable(qconf));
            ctx.setQueryName(queryName);
            return executeAsyncInternal(sessionHandle, ctx);
        } finally {
            release(sessionHandle);
        }
    }

    /**
     * Returns the query handle if the same query is already launched by the user
     * in the same session.
     *
     * @param query
     * @param sessionHandle
     * @param conf
     * @param queryName
     * @return
     */
    private QueryHandle checkForDuplicateQuery(String query, LensSessionHandle sessionHandle, Configuration conf,
            String queryName) {
        // Get all active queries of this session.
        List<QueryHandle> activeQueries = getSession(sessionHandle).getActiveQueries();
        synchronized (activeQueries) {
            for (QueryHandle handle : activeQueries) {
                QueryContext context = allQueries.get(handle);
                if (queryName.equals(context.getQueryName()) && query.equals(context.getUserQuery())
                        && Arrays.equals(UtilityMethods.generateHashOfWritable(conf), context.getQueryConfHash())) {
                    return handle;
                }
            }
        }
        return null;
    }

    /**
     * Creates the context.
     *
     * @param query    the query
     * @param userName the user name
     * @param conf     the conf
     * @param qconf    the qconf
     * @return the query context
     * @throws LensException the lens exception
     */
    protected QueryContext createContext(String query, String userName, LensConf conf, Configuration qconf,
            long timeOutMillis) throws LensException {
        QueryContext ctx = new QueryContext(query, userName, conf, qconf, drivers.values());
        ctx.setExecuteTimeoutMillis(timeOutMillis);
        return ctx;
    }

    /**
     * Creates the context.
     *
     * @param pctx     the pctx
     * @param userName the user name
     * @param conf     the conf
     * @param qconf    the qconf
     * @return the query context
     * @throws LensException the lens exception
     */
    protected QueryContext createContext(PreparedQueryContext pctx, String userName, LensConf conf,
            Configuration qconf, long timeOutMillis) throws LensException {
        QueryContext ctx = new QueryContext(pctx, userName, conf, qconf);
        ctx.setExecuteTimeoutMillis(timeOutMillis);
        return ctx;
    }

    /**
     * Execute async internal.
     *
     * @param sessionHandle the session handle
     * @param ctx           the ctx
     * @return the query handle
     * @throws LensException the lens exception
     */
    private QueryHandle executeAsyncInternal(LensSessionHandle sessionHandle, QueryContext ctx)
            throws LensException {

        ctx.setLensSessionIdentifier(sessionHandle.getPublicId().toString());
        rewriteAndSelect(ctx);
        return submitQuery(ctx);
    }

    private QueryHandle submitQuery(final QueryContext ctx) throws LensException {
        synchronized (ctx) {
            QueryStatus before = ctx.getStatus();
            ctx.setStatus(QueryStatus.getQueuedStatus());
            queuedQueries.add(ctx);
            log.info("Added to Queued Queries:{}", ctx.getQueryHandleString());
            allQueries.put(ctx.getQueryHandle(), ctx);
            // Add to session's active query list
            getSession(SESSION_MAP.get(ctx.getLensSessionIdentifier())).addToActiveQueries(ctx.getQueryHandle());
            fireStatusChangeEvent(ctx, ctx.getStatus(), before);
            log.info("Returning handle {}", ctx.getQueryHandle().getHandleId());
            return ctx.getQueryHandle();
        }
    }

    /*
     * (non-Javadoc)
     *
     * @see org.apache.lens.server.api.query.QueryExecutionService#updateQueryConf(org.apache.lens.api.LensSessionHandle,
     * org.apache.lens.api.query.QueryHandle, org.apache.lens.api.LensConf)
     */
    @Override
    public boolean updateQueryConf(LensSessionHandle sessionHandle, QueryHandle queryHandle, LensConf newconf)
            throws LensException {
        try {
            log.info("UpdateQueryConf: session:{} queryHandle: {}", sessionHandle, queryHandle);
            acquire(sessionHandle);
            QueryContext ctx = getUpdatedQueryContext(sessionHandle, queryHandle);
            if (ctx != null && (ctx.queued())) {
                ctx.updateConf(newconf.getProperties());
                // TODO COnf changed event tobe raised
                return true;
            } else {
                return false;
            }
        } finally {
            release(sessionHandle);
        }
    }

    /*
     * (non-Javadoc)
     *
     * @see org.apache.lens.server.api.query.QueryExecutionService#updateQueryConf(org.apache.lens.api.LensSessionHandle,
     * org.apache.lens.api.query.QueryPrepareHandle, org.apache.lens.api.LensConf)
     */
    @Override
    public boolean updateQueryConf(LensSessionHandle sessionHandle, QueryPrepareHandle prepareHandle,
            LensConf newconf) throws LensException {
        try {
            log.info("UpdatePreparedQueryConf: session:{} prepareHandle:{}", sessionHandle, prepareHandle);
            acquire(sessionHandle);
            PreparedQueryContext ctx = getPreparedQueryContext(sessionHandle, prepareHandle);
            ctx.updateConf(newconf.getProperties());
            return true;
        } finally {
            release(sessionHandle);
        }
    }

    /**
     *  Gets the query context either form memory or from DB (after query is purged)
     *  Note: For non-purged queries the status is updated before returning the context
     *
     * @param sessionHandle the session handle
     * @param queryHandle   the query handle
     * @return the query context
     * @throws LensException the lens exception
     */
    QueryContext getUpdatedQueryContext(LensSessionHandle sessionHandle, QueryHandle queryHandle)
            throws LensException {
        return getUpdatedQueryContext(sessionHandle, queryHandle, false);
    }

    /**
     * Gets the query context. If the query has been purged, null context is returned if returnNullIfPurged is true, else
     * context is read form DB Note: For non-purged queries the status is updated before returning the context
     *
     * @param sessionHandle
     * @param queryHandle
     * @param returnNullIfPurged
     * @return
     * @throws LensException
     */
    QueryContext getUpdatedQueryContext(LensSessionHandle sessionHandle, QueryHandle queryHandle,
            boolean returnNullIfPurged) throws LensException {
        try {
            acquire(sessionHandle);
            QueryContext ctx = allQueries.get(queryHandle);
            if (ctx == null) {
                return (returnNullIfPurged ? null : getQueryContextOfFinishedQuery(queryHandle));
            }
            updateStatus(queryHandle);
            return ctx;
        } finally {
            release(sessionHandle);
        }
    }

    QueryContext getQueryContextOfFinishedQuery(QueryHandle queryHandle) {
        FinishedLensQuery query = lensServerDao.getQuery(queryHandle.toString());
        log.info("FinishedLensQuery:{}", query);
        if (query == null) {
            throw new NotFoundException("Query not found " + queryHandle);
        }
        // pass the query conf instead of service conf
        return query.toQueryContext(conf, drivers.values());
    }

    /**
     * Gets the query context.
     * note: this method is made public to expose it to test cases
     * @param queryHandle the query handle
     * @return the query context
     */
    public QueryContext getQueryContext(QueryHandle queryHandle) {
        return allQueries.get(queryHandle);
    }

    /*
     * (non-Javadoc)
     *
     * @see org.apache.lens.server.api.query.QueryExecutionService#getQuery(org.apache.lens.api.LensSessionHandle,
     * org.apache.lens.api.query.QueryHandle)
     */
    @Override
    public LensQuery getQuery(LensSessionHandle sessionHandle, QueryHandle queryHandle) throws LensException {
        return getUpdatedQueryContext(sessionHandle, queryHandle).toLensQuery();
    }

    /**
     * Gets the prepared query context.
     *
     * @param sessionHandle the session handle
     * @param prepareHandle the prepare handle
     * @return the prepared query context
     * @throws LensException the lens exception
     */
    private PreparedQueryContext getPreparedQueryContext(LensSessionHandle sessionHandle,
            QueryPrepareHandle prepareHandle) throws LensException {
        try {
            acquire(sessionHandle);
            PreparedQueryContext ctx = preparedQueries.get(prepareHandle);
            if (ctx == null) {
                throw new NotFoundException("Prepared query not found " + prepareHandle);
            }
            return ctx;
        } finally {
            release(sessionHandle);
        }
    }

    /*
     * (non-Javadoc)
     *
     * @see org.apache.lens.server.api.query.QueryExecutionService#getPreparedQuery(org.apache.lens.api.LensSessionHandle,
     * org.apache.lens.api.query.QueryPrepareHandle)
     */
    @Override
    public LensPreparedQuery getPreparedQuery(LensSessionHandle sessionHandle, QueryPrepareHandle prepareHandle)
            throws LensException {
        return getPreparedQueryContext(sessionHandle, prepareHandle).toPreparedQuery();
    }

    /*
     * (non-Javadoc)
     *
     * @see org.apache.lens.server.api.query.QueryExecutionService#execute(org.apache.lens.api.LensSessionHandle,
     * java.lang.String, long, org.apache.lens.api.LensConf, java.lang.String)
     */
    @Override
    public QueryHandleWithResultSet execute(LensSessionHandle sessionHandle, String query, long timeoutMillis,
            LensConf conf, String queryName) throws LensException {
        try {
            log.info("Blocking execute session:{} query: {} timeout: {}", sessionHandle, query, timeoutMillis);
            acquire(sessionHandle);
            Configuration qconf = getLensConf(sessionHandle, conf);
            accept(query, qconf, SubmitOp.EXECUTE);
            QueryContext ctx = createContext(query, getSession(sessionHandle).getLoggedInUser(), conf, qconf,
                    timeoutMillis);
            ctx.setQueryName(queryName);
            ctx.setLensSessionIdentifier(sessionHandle.getPublicId().toString());
            rewriteAndSelect(ctx);
            return executeTimeoutInternal(sessionHandle, ctx, timeoutMillis, qconf);
        } finally {
            release(sessionHandle);
        }
    }

    /**
     * Execute timeout internal.
     *
     * @param sessionHandle the session handle
     * @param ctx           the ctx
     * @param timeoutMillis the timeout millis
     * @param conf          the conf
     * @return the query handle with result set
     * @throws LensException the lens exception
     */
    private QueryHandleWithResultSet executeTimeoutInternal(LensSessionHandle sessionHandle, QueryContext ctx,
            long timeoutMillis, Configuration conf) throws LensException {
        QueryHandle handle = submitQuery(ctx);
        long timeOutTime = ctx.getSubmissionTime() + timeoutMillis;
        log.info("query {} is executed with a timeout of {} and will timeout by {}", handle, timeoutMillis,
                timeOutTime);
        QueryHandleWithResultSet result = new QueryHandleWithResultSet(handle);

        boolean isQueued = true;
        while (isQueued && System.currentTimeMillis() < timeOutTime) {
            try {
                Thread.sleep(10);
                isQueued = isQueued(sessionHandle, handle);
            } catch (InterruptedException e) {
                log.error("Encountered Interrupted exception.", e);
            }
        }
        if (isQueued) { //query is still queued even after waiting for timeoutMillis
            result.setStatus(ctx.getStatus());
            addQueryToCancellationPool(ctx, conf, timeoutMillis); //cancel the timed-out Query
            return result;
        }

        QueryContext queryCtx = getUpdatedQueryContext(sessionHandle, handle);
        if (queryCtx.getSelectedDriver() == null) {
            result.setStatus(queryCtx.getStatus());
            return result;
        }

        QueryCompletionListenerImpl listener = new QueryCompletionListenerImpl(handle);
        long totalWaitTime = timeOutTime - System.currentTimeMillis();

        if (System.currentTimeMillis() < timeOutTime && !queryCtx.getStatus().executed()
                && !queryCtx.getStatus().finished()) {
            log.info("Registering for query {} completion notification", ctx.getQueryHandleString());
            queryCtx.getSelectedDriver().registerForCompletionNotification(ctx, totalWaitTime, listener);
            try {
                // We will wait for a few millis at a time until we reach max required wait time and also check the state
                // each time we come out of the wait.
                // This is done because the registerForCompletionNotification and query execution completion can happen
                // parallely especially in case of drivers like JDBC and in that case completion notification may not be
                // received by this listener. So its better to break the wait into smaller ones.
                long waitMillisPerCheck = totalWaitTime / 10;
                waitMillisPerCheck = (waitMillisPerCheck > 500) ? 500 : waitMillisPerCheck; // Lets keep max as 0.5 sec
                synchronized (listener) {
                    while (System.currentTimeMillis() < timeOutTime && !listener.querySuccessful
                            && !queryCtx.getStatus().executed() && !queryCtx.getStatus().finished()) {
                        listener.wait(waitMillisPerCheck);
                        if (!listener.querySuccessful) {
                            //update ths status in case query is not successful yet
                            queryCtx = getUpdatedQueryContext(sessionHandle, handle);
                        }
                    }
                }
            } catch (InterruptedException e) {
                log.info("{} query completion notification wait interrupted", queryCtx.getQueryHandleString());
            }
        }

        // At this stage (since the listener waits only for driver completion and not server that may include result
        // formatting and persistence) the query status can be RUNNING or EXECUTED or FAILED or SUCCESSFUL
        LensResultSet resultSet = null;
        queryCtx = getUpdatedQueryContext(sessionHandle, handle, true); // If the query is already purged queryCtx = null
        if (queryCtx != null && queryCtx.getStatus().isResultSetAvailable()) {
            resultSet = queryCtx.getSelectedDriver().fetchResultSet(queryCtx);
            if (resultSet instanceof PartiallyFetchedInMemoryResultSet) {
                PartiallyFetchedInMemoryResultSet partialnMemoryResult = (PartiallyFetchedInMemoryResultSet) resultSet;
                if (partialnMemoryResult.isComplteleyFetched()) { // DO not stream the result if its not completely fetched
                    result.setResult(new InMemoryQueryResult(partialnMemoryResult.getPreFetchedRows()));
                    result.setResultMetadata(partialnMemoryResult.getMetadata().toQueryResultSetMetadata());
                    result.setStatus(queryCtx.getStatus());
                    return result;
                }
            }
        }

        // Until timeOutTime, give this query a chance to reach FINISHED status if not already there.
        queryCtx = getUpdatedQueryContext(sessionHandle, handle);
        while (!queryCtx.finished() && System.currentTimeMillis() < timeOutTime) {
            queryCtx = getUpdatedQueryContext(sessionHandle, handle);
            try {
                Thread.sleep(100);
            } catch (InterruptedException e) {
                // Ignore
            }
        }

        if (queryCtx.finished() && queryCtx.getStatus().isResultSetAvailable()) {
            resultSet = getResultset(handle);
            result.setResultMetadata(resultSet.getMetadata().toQueryResultSetMetadata());
            result.setResult(resultSet.toQueryResult());
            result.setStatus(queryCtx.getStatus());
            return result;
        }

        // Result is not available. (Explicitly setting values to null for readability)
        result.setResult(null);
        result.setResultMetadata(null);
        result.setStatus(queryCtx.getStatus());

        if (!queryCtx.finished()) {
            addQueryToCancellationPool(queryCtx, conf, timeoutMillis); //cancel the timed-out Query
        }

        return result;
    }

    /**
     * This method is used to add a timed out query to cancellation pool.
     * The query gets cancelled asynchronously
     * Note : lens.query.cancel.on.timeout should be true for cancellation
     */
    private void addQueryToCancellationPool(QueryContext queryCtx, Configuration config, long timeoutMillis) {
        if (config.getBoolean(CANCEL_QUERY_ON_TIMEOUT, DEFAULT_CANCEL_QUERY_ON_TIMEOUT)) {
            log.info("Query {} could not be completed within the specified timeout interval. It will be cancelled",
                    queryCtx.getQueryHandleString());
            queryCancellationPool.submit(new CancelQueryTask(queryCtx.getQueryHandle()));
        } else {
            log.info(
                    "Query {} could not be completed within the specified timeout interval. Query cancellation is disabled",
                    queryCtx.getQueryHandleString());
        }
    }

    private boolean isQueued(final LensSessionHandle sessionHandle, final QueryHandle handle) throws LensException {
        // getQueryContext calls updateStatus, which fires query events if there's a change in status
        QueryContext query = getUpdatedQueryContext(sessionHandle, handle);
        synchronized (query) {
            return query.queued();
        }
    }

    /**
     * The Class QueryCompletionListenerImpl.
     */
    @Data
    class QueryCompletionListenerImpl extends QueryCompletionListener {

        /**
         * The succeeded.
         */
        boolean querySuccessful = false;

        /**
         * The handle.
         */
        final QueryHandle handle;

        @Override
        public void onCompletion(QueryHandle handle) {
            synchronized (this) {
                querySuccessful = true;
                log.info("Query {} with time out succeeded", handle);
                this.notify();
            }
        }

        @Override
        public void onError(QueryHandle handle, String error) {
            synchronized (this) {
                querySuccessful = false;
                log.info("Query {} with time out failed", handle);
                this.notify();
            }
        }
    }

    /*
     * (non-Javadoc)
     *
     * @see
     * org.apache.lens.server.api.query.QueryExecutionService#getResultSetMetadata(org.apache.lens.api.LensSessionHandle,
     * org.apache.lens.api.query.QueryHandle)
     */
    @Override
    public QueryResultSetMetadata getResultSetMetadata(LensSessionHandle sessionHandle, QueryHandle queryHandle)
            throws LensException {
        try {
            log.info("GetResultSetMetadata: session:{} query: {}", sessionHandle, queryHandle);
            acquire(sessionHandle);
            LensResultSet resultSet = getResultset(queryHandle);
            if (resultSet != null) {
                return resultSet.getMetadata().toQueryResultSetMetadata();
            } else {
                throw new NotFoundException(
                        "Resultset metadata not found for query: (" + sessionHandle + ", " + queryHandle + ")");
            }
        } finally {
            release(sessionHandle);
        }
    }

    /*
     * (non-Javadoc)
     *
     * @see org.apache.lens.server.api.query.QueryExecutionService#fetchResultSet(org.apache.lens.api.LensSessionHandle,
     * org.apache.lens.api.query.QueryHandle, long, int)
     */
    @Override
    public QueryResult fetchResultSet(LensSessionHandle sessionHandle, QueryHandle queryHandle, long startIndex,
            int fetchSize) throws LensException {
        try {
            log.info("FetchResultSet: session:{} query:{}", sessionHandle, queryHandle);
            acquire(sessionHandle);
            return getResultset(queryHandle).toQueryResult();
        } finally {
            release(sessionHandle);
        }
    }

    /*
     * (non-Javadoc)
     *
     * @see org.apache.lens.server.api.query.QueryExecutionService#closeResultSet(org.apache.lens.api.LensSessionHandle,
     * org.apache.lens.api.query.QueryHandle)
     */
    @Override
    public void closeResultSet(LensSessionHandle sessionHandle, QueryHandle queryHandle) throws LensException {
        try {
            log.info("CloseResultSet:session:{} query:{}", sessionHandle, queryHandle);
            acquire(sessionHandle);
            resultSets.remove(queryHandle);
            // Ask driver to close result set
            QueryContext ctx = getQueryContext(queryHandle);
            if (null != ctx) {
                ctx.getSelectedDriver().closeResultSet(queryHandle);
            }
        } finally {
            release(sessionHandle);
        }
    }

    /*
     * (non-Javadoc)
     *
     * @see org.apache.lens.server.api.query.QueryExecutionService#cancelQuery(org.apache.lens.api.LensSessionHandle,
     * org.apache.lens.api.query.QueryHandle)
     */
    @Override
    public boolean cancelQuery(LensSessionHandle sessionHandle, QueryHandle queryHandle) throws LensException {
        try {
            logSegregationContext.setLogSegragationAndQueryId(queryHandle.getHandleIdString());
            log.info("CancelQuery: session:{} query:{}", sessionHandle, queryHandle);
            acquire(sessionHandle);
            return cancelQuery(queryHandle);
        } finally {
            release(sessionHandle);
        }
    }

    private boolean cancelQuery(@NonNull QueryHandle queryHandle) throws LensException {
        QueryContext ctx = allQueries.get(queryHandle);
        if (ctx == null) {
            log.info("Could not cancel query {} as it has been purged already", queryHandle);
            return false;
        }

        synchronized (ctx) {

            updateStatus(queryHandle);

            if (ctx.finished()) {
                log.info("Could not cancel query {} as it has finished execution already", queryHandle);
                return false;
            }

            if (ctx.isLaunching()) {
                boolean launchCancelled = ctx.getQueryLauncher().cancel(true);
                log.info("query launch cancellation success : {}", launchCancelled);

            }
            if (ctx.launched() || ctx.running()) {
                if (!ctx.getSelectedDriver().cancelQuery(queryHandle)) {
                    log.info("Could not cancel query {}", queryHandle);
                    return false;
                }
            }

            log.info("Query {} cancelled successfully", queryHandle);
            setCancelledStatus(ctx, "Query is cancelled");
            return true;
        }
    }

    /*
     * (non-Javadoc)
     *
     * @see org.apache.lens.server.api.query.QueryExecutionService#getAllQueryDetails(
     * org.apache.lens.api.LensSessionHandle, java.lang.String, java.lang.String,
     * java.lang.String, java.lang.String, long, long)
     */
    @Override
    public List<LensQuery> getAllQueryDetails(LensSessionHandle sessionHandle, String states, String userName,
            String driver, String queryName, String fromDate, String toDate) throws LensException {
        long fromTime = -1;
        long toTime = Long.MAX_VALUE;
        Date now = new Date();
        if (fromDate != null) {
            fromTime = DateUtil.resolveDate(fromDate, now).getTime();
        }
        if (toDate != null) {
            toTime = DateUtil.resolveDate(toDate, now).getTime();
        }
        validateTimeRange(fromTime, toTime);
        try {
            acquire(sessionHandle);

            if (StringUtils.isBlank(userName)) {
                userName = getSession(sessionHandle).getLoggedInUser();
            }
            Set<Status> statuses = getStatuses(states);
            List<QueryHandle> inMemoryHandles = getQueriesInMemory(statuses, userName, driver, queryName, fromTime,
                    toTime);
            Set<LensQuery> result = new HashSet<>();
            for (QueryHandle handle : inMemoryHandles) {
                QueryContext ctx = allQueries.get(handle);
                if (ctx == null) {
                    continue;
                }
                result.add(ctx.toLensQuery());
            }

            List<LensQuery> persistedQueries = getPersistedQueries(userName, driver, queryName, fromTime, toTime,
                    statuses);
            result.addAll(persistedQueries);

            return new ArrayList<>(result);
        } finally {
            release(sessionHandle);
        }
    }

    /*
     * (non-Javadoc)
     *
     * @see org.apache.lens.server.api.query.QueryExecutionService#getAllQueries(org.apache.lens.api.LensSessionHandle,
     * java.lang.String, java.lang.String, java.lang.String, java.lang.String, long, long)
     */
    @Override
    public List<QueryHandle> getAllQueries(LensSessionHandle sessionHandle, String states, String userName,
            String driver, String queryName, String fromDate, String toDate) throws LensException {
        long fromTime = -1;
        long toTime = Long.MAX_VALUE;
        Date now = new Date();
        if (fromDate != null) {
            fromTime = DateUtil.resolveDate(fromDate, now).getTime();
        }
        if (toDate != null) {
            toTime = DateUtil.resolveDate(toDate, now).getTime();
        }
        validateTimeRange(fromTime, toTime);

        try {
            acquire(sessionHandle);

            if (StringUtils.isBlank(userName)) {
                userName = getSession(sessionHandle).getLoggedInUser();
            }
            Set<Status> statuses = getStatuses(states);
            List<QueryHandle> result = getQueriesInMemory(statuses, userName, driver, queryName, fromTime, toTime);
            List<QueryHandle> persistedQueries = getPersistedQueryHandles(userName, driver, queryName, fromTime,
                    toTime, statuses);

            HashSet<QueryHandle> deduplicatedResults = new HashSet<>(result);
            deduplicatedResults.addAll(persistedQueries);
            return new ArrayList<>(deduplicatedResults);
        } finally {
            release(sessionHandle);
        }
    }

    private List<LensQuery> getPersistedQueries(String userName, String driver, String queryName, long fromTime,
            long toTime, Set<Status> statuses) throws LensException {
        // Unless user wants to get queries in 'non finished' state, get finished queries from DB as well
        List<Status> finishedStatusesQueried = finishedStatuses(statuses);
        if (finishedStatusesQueried == null || !finishedStatusesQueried.isEmpty()) {
            if ("all".equalsIgnoreCase(userName)) {
                userName = null;
            }
            List<FinishedLensQuery> results = lensServerDao.findFinishedQueryDetails(finishedStatusesQueried,
                    userName, driver, queryName, fromTime, toTime);
            List<LensQuery> persistedQueries = new ArrayList<>();
            for (FinishedLensQuery finishedLensQuery : results) {
                persistedQueries.add(finishedLensQuery.toQueryContext(conf, drivers.values()).toLensQuery());
            }
            if (!persistedQueries.isEmpty()) {
                log.info("Adding persisted queries {}", persistedQueries.size());
                return persistedQueries;
            }
        }
        return Collections.emptyList();
    }

    private List<Status> finishedStatuses(Set<Status> statuses) {
        List<Status> finishedStatusesQueried = null;
        if (statuses.size() != Status.values().length) {
            finishedStatusesQueried = Lists.newArrayList();
            for (Status status : statuses) {
                switch (status) {
                case CANCELED:
                case SUCCESSFUL:
                case FAILED:
                    finishedStatusesQueried.add(status);
                    break;
                default:
                    break;
                }
            }
        }
        return finishedStatusesQueried;
    }

    private List<QueryHandle> getPersistedQueryHandles(String userName, String driver, String queryName,
            long fromTime, long toTime, Set<Status> statuses) throws LensException {
        // Unless user wants to get queries in 'non finished' state, get finished queries from DB as well
        List<Status> finishedStatusesQueried = finishedStatuses(statuses);
        if (finishedStatusesQueried == null || !finishedStatusesQueried.isEmpty()) {
            if ("all".equalsIgnoreCase(userName)) {
                userName = null;
            }
            List<QueryHandle> persistedQueries = lensServerDao.findFinishedQueries(finishedStatusesQueried,
                    userName, driver, queryName, fromTime, toTime);
            if (persistedQueries != null && !persistedQueries.isEmpty()) {
                log.info("Adding persisted queries {}", persistedQueries.size());
                return persistedQueries;
            }
        }
        return Collections.emptyList();
    }

    private List<QueryHandle> getQueriesInMemory(Set<Status> statuses, String userName, String driver,
            String queryName, long fromTime, long toTime) throws LensException {

        userName = UtilityMethods.removeDomain(userName);
        boolean filterByQueryName = StringUtils.isNotBlank(queryName);
        if (filterByQueryName) {
            queryName = queryName.toLowerCase();
        }

        boolean filterByDriver = StringUtils.isNotBlank(driver);

        List<QueryHandle> all = new ArrayList<QueryHandle>(allQueries.keySet());
        Iterator<QueryHandle> itr = all.iterator();
        while (itr.hasNext()) {
            QueryHandle q = itr.next();
            QueryContext context = allQueries.get(q);
            long querySubmitTime = context.getSubmissionTime();
            if ((!statuses.contains(context.getStatus().getStatus()))
                    || (filterByQueryName && !context.getQueryName().toLowerCase().contains(queryName))
                    || (filterByDriver
                            && !context.getSelectedDriver().getFullyQualifiedName().equalsIgnoreCase(driver))
                    || (!"all".equalsIgnoreCase(userName) && !userName.equalsIgnoreCase(context.getSubmittedUser()))
                    || (!(fromTime <= querySubmitTime && querySubmitTime < toTime))) {
                itr.remove();
            }
        }
        return all;
    }

    private Set<Status> getStatuses(String states) {
        Set<Status> statuses;
        try {
            if (StringUtils.isNotBlank(states)) {
                statuses = Sets.newHashSet();
                for (String state : states.split(",")) {
                    statuses.add(Status.valueOf(state.trim().toUpperCase()));
                }
            } else {
                statuses = Sets.newHashSet(Status.values());
            }
        } catch (IllegalArgumentException e) {
            throw new BadRequestException(
                    "Bad state argument passed, possible values are " + Arrays.toString(Status.values()), e);
        }
        return statuses;
    }

    /*
     * (non-Javadoc)
     *
     * @see
     * org.apache.lens.server.api.query.QueryExecutionService#getAllPreparedQueries(org.apache.lens.api.LensSessionHandle,
     * java.lang.String, java.lang.String, long, long)
     */
    @Override
    public List<QueryPrepareHandle> getAllPreparedQueries(LensSessionHandle sessionHandle, String user,
            String queryName, String fromDate, String toDate) throws LensException {
        long fromTime = -1;
        long toTime = Long.MAX_VALUE;
        Date now = new Date();
        if (fromDate != null) {
            fromTime = DateUtil.resolveDate(fromDate, now).getTime();
        }
        if (toDate != null) {
            toTime = DateUtil.resolveDate(toDate, now).getTime();
        }
        validateTimeRange(fromTime, toTime);
        user = UtilityMethods.removeDomain(user);
        try {
            acquire(sessionHandle);
            List<QueryPrepareHandle> allPrepared = new ArrayList<QueryPrepareHandle>(preparedQueries.keySet());
            Iterator<QueryPrepareHandle> itr = allPrepared.iterator();
            while (itr.hasNext()) {
                QueryPrepareHandle q = itr.next();
                PreparedQueryContext preparedQueryContext = preparedQueries.get(q);

                if (StringUtils.isNotBlank(user)) {
                    if ("all".equalsIgnoreCase(user)) {
                        continue;
                    } else if (user.equalsIgnoreCase(preparedQueryContext.getPreparedUser())) {
                        continue;
                    }
                }

                if (StringUtils.isNotBlank(queryName)) {
                    if (preparedQueryContext.getQueryName().toLowerCase().contains(queryName.toLowerCase())) {
                        continue;
                    }
                }
                long queryPrepTime = preparedQueryContext.getPreparedTime().getTime();
                if (fromTime <= queryPrepTime && queryPrepTime < toTime) {
                    continue;
                }
                itr.remove();
            }
            return allPrepared;
        } finally {
            release(sessionHandle);
        }
    }

    /**
     * Validate time range.
     *
     * @param fromDate the from date
     * @param toDate   the to date
     */
    private void validateTimeRange(long fromDate, long toDate) {
        if (fromDate >= toDate) {
            throw new BadRequestException("Invalid time range: [" + fromDate + ", " + toDate + "]");
        }
    }

    /*
     * (non-Javadoc)
     *
     * @see org.apache.lens.server.api.query.QueryExecutionService#destroyPrepared(org.apache.lens.api.LensSessionHandle,
     * org.apache.lens.api.query.QueryPrepareHandle)
     */
    @Override
    public boolean destroyPrepared(LensSessionHandle sessionHandle, QueryPrepareHandle prepared)
            throws LensException {
        try {
            log.info("DestroyPrepared: {} prepareHandle:{}", sessionHandle, prepared);
            acquire(sessionHandle);
            destroyPreparedQuery(getPreparedQueryContext(sessionHandle, prepared));
            return true;
        } finally {
            release(sessionHandle);
        }
    }

    /**
     * Destroy prepared query.
     *
     * @param ctx the ctx
     * @throws LensException the lens exception
     */
    private void destroyPreparedQuery(PreparedQueryContext ctx) throws LensException {
        if (ctx.getSelectedDriver() != null) {
            ctx.getSelectedDriver().closePreparedQuery(ctx.getPrepareHandle());
        }
        preparedQueries.remove(ctx.getPrepareHandle());
        preparedQueryQueue.remove(ctx);
        decrCounter(PREPARED_QUERIES_COUNTER);
    }

    /*
     * (non-Javadoc)
     *
     * @see org.apache.lens.server.api.query.QueryExecutionService#estimate(org.apache.lens.api.LensSessionHandle,
     * java.lang.String, org.apache.lens.api.LensConf)
     */
    @Override
    public QueryCost estimate(final String requestId, LensSessionHandle sessionHandle, String query,
            LensConf lensConf) throws LensException {
        try {
            log.info("Estimate: session :{} query:{}", sessionHandle, query);
            acquire(sessionHandle);
            Configuration qconf = getLensConf(sessionHandle, lensConf);
            ExplainQueryContext estimateQueryContext = new ExplainQueryContext(requestId, query,
                    getSession(sessionHandle).getLoggedInUser(), lensConf, qconf, drivers.values());
            estimateQueryContext.setLensSessionIdentifier(sessionHandle.getPublicId().toString());
            accept(query, qconf, SubmitOp.ESTIMATE);
            rewriteAndSelect(estimateQueryContext);
            return estimateQueryContext.getSelectedDriverQueryCost();
        } finally {
            release(sessionHandle);
        }
    }

    /*
     * (non-Javadoc)
     *
     * @see org.apache.lens.server.api.query.QueryExecutionService#explain(org.apache.lens.api.LensSessionHandle,
     * java.lang.String, org.apache.lens.api.LensConf)
     */
    @Override
    public QueryPlan explain(final String requestId, LensSessionHandle sessionHandle, String query,
            LensConf lensConf) throws LensException {
        try {
            log.info("Explain: session:{} query:{}", sessionHandle, query);
            acquire(sessionHandle);
            Configuration qconf = getLensConf(sessionHandle, lensConf);
            ExplainQueryContext explainQueryContext = new ExplainQueryContext(requestId, query,
                    getSession(sessionHandle).getLoggedInUser(), lensConf, qconf, drivers.values());
            explainQueryContext.setLensSessionIdentifier(sessionHandle.getPublicId().toString());
            accept(query, qconf, SubmitOp.EXPLAIN);
            rewriteAndSelect(explainQueryContext);
            addSessionResourcesToDriver(explainQueryContext);
            return explainQueryContext.getSelectedDriver().explain(explainQueryContext).toQueryPlan();
        } catch (UnsupportedEncodingException e) {
            throw new LensException(e);
        } finally {
            release(sessionHandle);
        }
    }

    /**
     * Creates the add/delete resource query.
     *
     * @param command
     * @param sessionHandle
     * @param driver
     * @return
     * @throws LensException
     */
    private QueryContext createResourceQuery(String command, LensSessionHandle sessionHandle, LensDriver driver)
            throws LensException {
        LensConf qconf = new LensConf();
        qconf.addProperty(QUERY_PERSISTENT_RESULT_INDRIVER, "false");
        QueryContext addQuery = QueryContext.createContextWithSingleDriver(command,
                getSession(sessionHandle).getLoggedInUser(), qconf, getLensConf(sessionHandle, qconf), driver,
                sessionHandle.getPublicId().toString(), true);
        return addQuery;
    }

    /*
     * (non-Javadoc)
     *
     * @see org.apache.lens.server.LensService#deleteResource(org.apache.lens.api.LensSessionHandle, java.lang.String,
     * java.lang.String)
     */
    public void deleteResource(LensSessionHandle sessionHandle, String type, String path) throws LensException {
        try {
            acquire(sessionHandle);
            String command = "delete " + type.toLowerCase() + " " + path;
            for (LensDriver driver : drivers.values()) {
                if (driver instanceof HiveDriver) {
                    driver.execute(createResourceQuery(command, sessionHandle, driver));
                }
            }
        } finally {
            release(sessionHandle);
        }
    }

    /*
     * (non-Javadoc)
     *
     * @see org.apache.lens.server.LensService#readExternal(java.io.ObjectInput)
     */
    @Override
    public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException {
        super.readExternal(in);
        // Restore drivers
        synchronized (drivers) {
            int numDrivers = in.readInt();
            String driverQualifiedName;
            String driverClsName;
            for (int i = 0; i < numDrivers; i++) {
                driverQualifiedName = in.readUTF();
                driverClsName = in.readUTF();
                LensDriver driver = drivers.get(driverQualifiedName);
                if (driver == null) {
                    // this driver is removed in the current server restart
                    // we will create an instance and read its state still.
                    try {
                        Class<? extends LensDriver> driverCls = (Class<? extends LensDriver>) Class
                                .forName(driverClsName);
                        driver = driverCls.newInstance();
                        String[] driverTypeAndName = StringUtils.split(driverQualifiedName, '/');
                        driver.configure(conf, driverTypeAndName[0], driverTypeAndName[1]);
                    } catch (Exception e) {
                        log.error("Could not instantiate driver:{} represented by class {}", driverQualifiedName,
                                driverClsName, e);
                        throw new IOException(e);
                    }
                    log.info("Driver state for {} will be ignored", driverQualifiedName);
                }
                driver.readExternal(in);
            }
        }

        // Restore queries
        synchronized (allQueries) {
            int numQueries = in.readInt();

            for (int i = 0; i < numQueries; i++) {
                QueryContext ctx = (QueryContext) in.readObject();
                ctx.initTransientState();

                //Create DriverSelectorQueryContext by passing all the drivers and the user query
                //Driver conf gets reset in start
                DriverSelectorQueryContext driverCtx = new DriverSelectorQueryContext(ctx.getUserQuery(),
                        new Configuration(), drivers.values());
                ctx.setDriverContext(driverCtx);
                boolean driverAvailable = in.readBoolean();
                // set the selected driver if available, if not available for the cases of queued queries,
                // query service will do the selection from existing drivers and update
                if (driverAvailable) {
                    String selectedDriverQualifiedName = in.readUTF();
                    if (drivers.get(selectedDriverQualifiedName) != null) {
                        ctx.getDriverContext().setSelectedDriver(drivers.get(selectedDriverQualifiedName));
                        ctx.setDriverQuery(ctx.getSelectedDriver(), ctx.getSelectedDriverQuery());
                    } else {
                        log.info("Ignoring {} as the driver is not loaded", ctx.getQueryHandle());
                        continue;
                    }
                }
                allQueries.put(ctx.getQueryHandle(), ctx);
            }
        }
    }

    /*
     * (non-Javadoc)
     *
     * @see org.apache.lens.server.LensService#writeExternal(java.io.ObjectOutput)
     */
    @Override
    public void writeExternal(ObjectOutput out) throws IOException {
        super.writeExternal(out);
        // persist all drivers
        synchronized (drivers) {
            out.writeInt(drivers.size());
            LensDriver driver = null;
            for (Map.Entry<String, LensDriver> driverEntry : drivers.entrySet()) {
                driver = driverEntry.getValue();
                synchronized (driver) {
                    out.writeUTF(driverEntry.getKey());
                    out.writeUTF(driver.getClass().getName());
                    driver.writeExternal(out);
                }
            }
        }
        // persist allQueries
        synchronized (allQueries) {
            out.writeInt(allQueries.size());
            for (QueryContext ctx : allQueries.values()) {
                synchronized (ctx) {
                    out.writeObject(ctx);
                    boolean isDriverAvailable = (ctx.getSelectedDriver() != null);
                    out.writeBoolean(isDriverAvailable);
                    if (isDriverAvailable) {
                        out.writeUTF(ctx.getSelectedDriver().getFullyQualifiedName());
                    }
                }
            }
            log.info("Persisted {} queries", allQueries.size());
        }
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public HealthStatus getHealthStatus() {
        boolean isHealthy = true;
        StringBuilder details = new StringBuilder();

        if (!this.getServiceState().equals(STATE.STARTED)) {
            isHealthy = false;
            details.append("Query execution service is down.");
        }

        if (!this.statusPoller.isAlive()) {
            isHealthy = false;
            details.append("Status poller thread is dead.");
        }

        if (!this.prepareQueryPurger.isAlive()) {
            isHealthy = false;
            details.append("PrepareQuery purger thread is dead.");
        }

        if (!this.queryPurger.isAlive()) {
            isHealthy = false;
            details.append("Query purger thread is dead.");
        }

        if (!this.querySubmitter.isAlive()) {
            isHealthy = false;
            details.append("Query submitter thread is dead.");
        }

        if (this.queryLauncherPool.isShutdown() || this.queryLauncherPool.isTerminated()) {
            isHealthy = false;
            details.append("Query launcher Pool is dead.");
        }

        if (this.estimatePool.isShutdown() || this.estimatePool.isTerminated()) {
            isHealthy = false;
            details.append("Estimate Pool is dead.");
        }

        if (querySubmitterRunnable.pausedForTest) {
            isHealthy = false;
            details.append("QuerySubmitter paused for test.");
        }

        if (null != this.queryResultPurger && !this.queryResultPurger.isHealthy()) {
            isHealthy = false;
            details.append("QueryResultPurger is dead.");
        }

        if (queryCancellationPool.isShutdown() || queryCancellationPool.isTerminated()) {
            isHealthy = false;
            details.append("Query Cancellation Pool is dead.");
        }

        if (queryExpirer.isShutdown() || queryExpirer.isTerminated()) {
            isHealthy = false;
            details.append("Query Expiry thread is dead.");
        }

        if (!isHealthy) {
            log.error(details.toString());
        }

        return isHealthy ? new HealthStatus(isHealthy, "QueryExecution service is healthy.")
                : new HealthStatus(isHealthy, details.toString());
    }
    /*
     * (non-Javadoc)
     *
     * @see org.apache.lens.server.api.query.QueryExecutionService#getHttpResultSet(org.apache.lens.api.LensSessionHandle,
     * org.apache.lens.api.query.QueryHandle)
     */

    @Override
    public Response getHttpResultSet(LensSessionHandle sessionHandle, QueryHandle queryHandle)
            throws LensException {
        LensResultSet resultSet = getResultset(queryHandle);
        if (!resultSet.isHttpResultAvailable()) {
            throw new NotFoundException("http result not available");
        }
        final Path resultPath = new Path(resultSet.getOutputPath());
        try {
            FileSystem fs = resultPath.getFileSystem(conf);
            if (!fs.exists(resultPath)) {
                throw new NotFoundException("Result file does not exist!");
            }
        } catch (IOException e) {
            throw new LensException(e);
        }
        final QueryContext ctx = getUpdatedQueryContext(sessionHandle, queryHandle);
        String resultFSReadUrl = conf.get(RESULT_FS_READ_URL);
        if (resultFSReadUrl != null) {
            try {
                URI resultReadPath = new URI(resultFSReadUrl + resultPath.toUri().getPath() + "?op=OPEN&user.name="
                        + getSession(sessionHandle).getClusterUser());
                return Response.seeOther(resultReadPath)
                        .header("content-disposition", "attachment; filename = " + resultPath.getName())
                        .type(MediaType.APPLICATION_OCTET_STREAM).build();
            } catch (URISyntaxException e) {
                throw new LensException(e);
            }
        } else {
            StreamingOutput stream = new StreamingOutput() {
                @Override
                public void write(OutputStream os) throws IOException {
                    FSDataInputStream fin = null;
                    try {
                        FileSystem fs = resultPath.getFileSystem(ctx.getConf());
                        fin = fs.open(resultPath);
                        UtilityMethods.pipe(fin, os);
                    } finally {
                        if (fin != null) {
                            fin.close();
                        }

                    }
                }
            };
            return Response.ok(stream)
                    .header("content-disposition", "attachment; filename = " + resultPath.getName())
                    .type(MediaType.APPLICATION_OCTET_STREAM).build();
        }
    }

    /**
     * Allow drivers to release resources acquired for a session if any.
     *
     * @param sessionHandle the session handle
     */
    public void closeDriverSessions(LensSessionHandle sessionHandle) {
        for (LensDriver driver : drivers.values()) {
            if (driver instanceof HiveDriver) {
                ((HiveDriver) driver).closeSession(sessionHandle);
            }
        }
    }

    // Used in test code
    Collection<LensDriver> getDrivers() {
        return drivers.values();
    }

    @Override
    public long getQueuedQueriesCount() {
        return queuedQueries.size();
    }

    @Override
    public long getRunningQueriesCount() {
        return launchedQueries.getQueriesCount();
    }

    @Override
    public long getWaitingQueriesCount() {
        return waitingQueries.getQueriesCount();
    }

    @Override
    public long getFinishedQueriesCount() {
        return finishedQueries.size();
    }

    @Override
    public long getLaunchingQueriesCount() {
        long count = 0;
        for (QueryContext ctx : launchedQueries.getQueries()) {
            if (ctx.isLaunching()) {
                count++;
            }
        }
        return count;
    }

    @Data
    @AllArgsConstructor
    public static class QueryCount {
        long running, queued, waiting;
    }

    public QueryCount getQueryCountSnapshot() {
        removalFromLaunchedQueriesLock.lock();
        QueryCount count = new QueryCount(getRunningQueriesCount(), getQueuedQueriesCount(),
                getWaitingQueriesCount());
        removalFromLaunchedQueriesLock.unlock();
        return count;
    }

    /**
     * Handle driver session start.
     *
     * @param event the event
     */
    protected void handleDriverSessionStart(DriverEvent event) {
        DriverSessionStarted sessionStarted = (DriverSessionStarted) event;
        if (!(event.getDriver() instanceof HiveDriver)) {
            return;
        }

        HiveDriver hiveDriver = (HiveDriver) event.getDriver();

        String lensSession = sessionStarted.getLensSessionID();
        LensSessionHandle sessionHandle = getSessionHandle(lensSession);
        if (sessionHandle == null) {
            log.warn("Lens session went away for sessionid:" + lensSession);
            return;
        }

        try {
            LensSessionImpl session = getSession(sessionHandle);
            acquire(sessionHandle);
            // Add resources for this session
            List<ResourceEntry> resources = session.getLensSessionPersistInfo().getResources();
            if (resources != null && !resources.isEmpty()) {
                for (ResourceEntry resource : resources) {
                    log.info("{} Restoring resource {} for session {}", hiveDriver, resource, lensSession);
                    String command = "add " + resource.getType().toLowerCase() + " " + resource.getUri();
                    try {
                        // Execute add resource query in blocking mode
                        hiveDriver.execute(createResourceQuery(command, sessionHandle, hiveDriver));
                        resource.restoredResource();
                        log.info("{} Restored resource {} for session {}", hiveDriver, resource, lensSession);
                    } catch (Exception exc) {
                        log.error("{} Unable to add resource {} for session {}", hiveDriver, resource, lensSession,
                                exc);
                    }
                }
            } else {
                log.info("{} No resources to restore for session {}", hiveDriver, lensSession);
            }
        } catch (Exception e) {
            log.warn("Lens session went away! {} driver session: {}", lensSession,
                    ((DriverSessionStarted) event).getDriverSessionID(), e);
        } finally {
            release(sessionHandle);
        }
    }

    /**
     * Add session's resources to selected driver if needed
     *
     * @param ctx QueryContext for executinf queries
     */
    protected void addSessionResourcesToDriver(final AbstractQueryContext ctx) {
        LensDriver driver = ctx.getSelectedDriver();
        String sessionIdentifier = ctx.getLensSessionIdentifier();

        if (!(driver instanceof HiveDriver) || StringUtils.isBlank(sessionIdentifier)) {
            // Adding resources only required for Hive driver
            return;
        }

        HiveDriver hiveDriver = (HiveDriver) driver;

        // Check if jars need to be passed to selected driver
        final LensSessionHandle sessionHandle = getSessionHandle(sessionIdentifier);
        final LensSessionImpl session = getSession(sessionHandle);

        // Add resources if either they haven't been marked as added on the session, or if Hive driver says they need
        // to be added to the corresponding hive driver
        if (!hiveDriver.areDBResourcesAddedForSession(sessionIdentifier, ctx.getDatabase())) {
            Collection<ResourceEntry> dbResources = session.getDBResources(ctx.getDatabase());

            if (CollectionUtils.isNotEmpty(dbResources)) {
                log.info("Proceeding to add resources for DB {} for query {} resources: {}",
                        session.getCurrentDatabase(), ctx.getLogHandle(), dbResources);

                List<ResourceEntry> failedDBResources = addResources(dbResources, sessionHandle, hiveDriver);
                Iterator<ResourceEntry> itr = dbResources.iterator();
                while (itr.hasNext()) {
                    ResourceEntry res = itr.next();
                    if (!failedDBResources.contains(res)) {
                        itr.remove();
                    }
                }
            } else {
                log.info("No need to add DB resources for session: {} db= {}", sessionIdentifier,
                        session.getCurrentDatabase());
            }
            hiveDriver.setResourcesAddedForSession(sessionIdentifier, ctx.getDatabase());
        }

        // Get pending session resources which needed to be added for this database
        Collection<ResourceEntry> pendingResources = session
                .getPendingSessionResourcesForDatabase(ctx.getDatabase());
        log.info("Adding pending {} session resources for session {} for database {}", pendingResources.size(),
                sessionIdentifier, ctx.getDatabase());
        List<ResourceEntry> failedResources = addResources(pendingResources, sessionHandle, hiveDriver);
        // Mark added resources so that we don't add them again. If any of the resources failed
        // to be added, then they will be added again
        for (ResourceEntry res : pendingResources) {
            if (!failedResources.contains(res)) {
                res.addToDatabase(ctx.getDatabase());
            }
        }
    }

    /**
     * Add resources to hive driver, returning resources which failed to be added
     *
     * @param resources     collection of resources intented to be added to hive driver
     * @param sessionHandle
     * @param hiveDriver
     * @return resources which could not be added to hive driver
     */
    private List<ResourceEntry> addResources(Collection<ResourceEntry> resources, LensSessionHandle sessionHandle,
            HiveDriver hiveDriver) {
        List<ResourceEntry> failedResources = new ArrayList<ResourceEntry>();
        for (ResourceEntry res : resources) {
            try {
                addSingleResourceToHive(hiveDriver, res, sessionHandle);
            } catch (LensException exc) {
                failedResources.add(res);
                log.error("Error adding resources for session {} resources: {}", sessionHandle, res.getUri(), exc);
            }
        }
        return failedResources;
    }

    private void addSingleResourceToHive(HiveDriver driver, ResourceEntry res, LensSessionHandle sessionHandle)
            throws LensException {
        String sessionIdentifier = sessionHandle.getPublicId().toString();
        String uri = res.getUri();
        // Hive doesn't and URIs starting with file:/ correctly, so we have to change it to file:///
        // See: org.apache.hadoop.hive.ql.exec.Utilities.addToClassPath
        uri = removePrefixBeforeURI(uri);

        String command = "add " + res.getType().toLowerCase() + " " + uri;
        driver.execute(createResourceQuery(command, sessionHandle, driver));
        log.info("Added resource to hive driver {} for session {} cmd: {}", driver, sessionIdentifier, command);
    }

    private boolean removeFromLaunchedQueries(final QueryContext finishedQuery) {

        /* Check javadoc of QueryExecutionServiceImpl#removalFromLaunchedQueriesLock for reason for existence
        of this lock. */

        log.debug("Acquiring lock in removeFromLaunchedQueries");
        removalFromLaunchedQueriesLock.lock();
        boolean modified = false;

        try {
            modified = this.launchedQueries.remove(finishedQuery);
        } finally {
            removalFromLaunchedQueriesLock.unlock();
        }

        log.debug("launchedQueries.remove(finishedQuery) has returned [{}] for finished query with query id:[{}]",
                modified, finishedQuery.getQueryHandleString());
        return modified;
    }

    /**
     * Caller of this method must make sure that this method is called inside a synchronized(queryContext) block
     * for a safe copy from queryContext to an instance of {@link FinishedLensQuery}
     *
     * @param queryContext
     */
    private void processWaitingQueriesAsync(final QueryContext queryContext) {

        final FinishedLensQuery finishedLensQuery = new FinishedLensQuery(queryContext);

        Runnable r = new Runnable() {
            @Override
            public void run() {
                try {
                    processWaitingQueries(finishedLensQuery);
                } catch (final Throwable e) {
                    log.error("Error in processing waiting queries", e);
                }
            }
        };

        exceptionSafeSubmit(this.waitingQueriesSelectionSvc, r);
    }

    private void exceptionSafeSubmit(final ExecutorService svc, final Runnable r) {
        try {
            svc.submit(r);
        } catch (final Throwable e) {
            log.debug("Could not submit runnable:{}", e);
        }
    }

    private void processWaitingQueries(final FinishedLensQuery finishedQuery) {

        Set<QueryContext> eligibleWaitingQueries = this.waitingQueriesSelector.selectQueries(finishedQuery,
                this.waitingQueries);
        log.info("Eligible queries to pick from waiting queries: {}", eligibleWaitingQueries);
        if (eligibleWaitingQueries.isEmpty()) {
            log.info("No queries eligible to move out of waiting state.");
            return;
        }
        waitingQueries.removeAll(eligibleWaitingQueries);
        queuedQueries.addAll(eligibleWaitingQueries);
        if (log.isDebugEnabled()) {
            log.debug("Added {} queries to queued queries", eligibleWaitingQueries.size());
        }
    }
}