org.apache.hive.service.cli.session.SessionManager.java Source code

Java tutorial

Introduction

Here is the source code for org.apache.hive.service.cli.session.SessionManager.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.hive.service.cli.session;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Date;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Future;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

import org.apache.commons.io.FileUtils;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.hadoop.hive.conf.HiveConf;
import org.apache.hadoop.hive.conf.HiveConf.ConfVars;
import org.apache.hive.service.CompositeService;
import org.apache.hive.service.cli.HiveSQLException;
import org.apache.hive.service.cli.SessionHandle;
import org.apache.hive.service.cli.operation.OperationManager;
import org.apache.hive.service.cli.thrift.TProtocolVersion;
import org.apache.hive.service.server.HiveServer2;
import org.apache.hive.service.server.ThreadFactoryWithGarbageCleanup;

/**
 * SessionManager.
 *
 */
public class SessionManager extends CompositeService {

    private static final Log LOG = LogFactory.getLog(CompositeService.class);
    public static final String HIVERCFILE = ".hiverc";
    private HiveConf hiveConf;
    private final Map<SessionHandle, HiveSession> handleToSession = new ConcurrentHashMap<SessionHandle, HiveSession>();
    private final OperationManager operationManager = new OperationManager();
    private ThreadPoolExecutor backgroundOperationPool;
    private boolean isOperationLogEnabled;
    private File operationLogRootDir;

    private long checkInterval;
    private long sessionTimeout;
    private boolean checkOperation;

    private volatile boolean shutdown;
    // The HiveServer2 instance running this service
    private final HiveServer2 hiveServer2;

    public SessionManager(HiveServer2 hiveServer2) {
        super(SessionManager.class.getSimpleName());
        this.hiveServer2 = hiveServer2;
    }

    @Override
    public synchronized void init(HiveConf hiveConf) {
        this.hiveConf = hiveConf;
        //Create operation log root directory, if operation logging is enabled
        if (hiveConf.getBoolVar(ConfVars.HIVE_SERVER2_LOGGING_OPERATION_ENABLED)) {
            initOperationLogRootDir();
        }
        createBackgroundOperationPool();
        addService(operationManager);
        super.init(hiveConf);
    }

    private void createBackgroundOperationPool() {
        int poolSize = hiveConf.getIntVar(ConfVars.HIVE_SERVER2_ASYNC_EXEC_THREADS);
        LOG.info("HiveServer2: Background operation thread pool size: " + poolSize);
        int poolQueueSize = hiveConf.getIntVar(ConfVars.HIVE_SERVER2_ASYNC_EXEC_WAIT_QUEUE_SIZE);
        LOG.info("HiveServer2: Background operation thread wait queue size: " + poolQueueSize);
        long keepAliveTime = HiveConf.getTimeVar(hiveConf, ConfVars.HIVE_SERVER2_ASYNC_EXEC_KEEPALIVE_TIME,
                TimeUnit.SECONDS);
        LOG.info("HiveServer2: Background operation thread keepalive time: " + keepAliveTime + " seconds");

        // Create a thread pool with #poolSize threads
        // Threads terminate when they are idle for more than the keepAliveTime
        // A bounded blocking queue is used to queue incoming operations, if #operations > poolSize
        String threadPoolName = "HiveServer2-Background-Pool";
        backgroundOperationPool = new ThreadPoolExecutor(poolSize, poolSize, keepAliveTime, TimeUnit.SECONDS,
                new LinkedBlockingQueue<Runnable>(poolQueueSize),
                new ThreadFactoryWithGarbageCleanup(threadPoolName));
        backgroundOperationPool.allowCoreThreadTimeOut(true);

        checkInterval = HiveConf.getTimeVar(hiveConf, ConfVars.HIVE_SERVER2_SESSION_CHECK_INTERVAL,
                TimeUnit.MILLISECONDS);
        sessionTimeout = HiveConf.getTimeVar(hiveConf, ConfVars.HIVE_SERVER2_IDLE_SESSION_TIMEOUT,
                TimeUnit.MILLISECONDS);
        checkOperation = HiveConf.getBoolVar(hiveConf, ConfVars.HIVE_SERVER2_IDLE_SESSION_CHECK_OPERATION);
    }

    private void initOperationLogRootDir() {
        operationLogRootDir = new File(hiveConf.getVar(ConfVars.HIVE_SERVER2_LOGGING_OPERATION_LOG_LOCATION));
        isOperationLogEnabled = true;

        if (operationLogRootDir.exists() && !operationLogRootDir.isDirectory()) {
            LOG.warn("The operation log root directory exists, but it is not a directory: "
                    + operationLogRootDir.getAbsolutePath());
            isOperationLogEnabled = false;
        }

        if (!operationLogRootDir.exists()) {
            if (!operationLogRootDir.mkdirs()) {
                LOG.warn("Unable to create operation log root directory: " + operationLogRootDir.getAbsolutePath());
                isOperationLogEnabled = false;
            }
        }

        if (isOperationLogEnabled) {
            LOG.info("Operation log root directory is created: " + operationLogRootDir.getAbsolutePath());
            try {
                FileUtils.forceDeleteOnExit(operationLogRootDir);
            } catch (IOException e) {
                LOG.warn("Failed to schedule cleanup HS2 operation logging root dir: "
                        + operationLogRootDir.getAbsolutePath(), e);
            }
        }
    }

    @Override
    public synchronized void start() {
        super.start();
        if (checkInterval > 0) {
            startTimeoutChecker();
        }
    }

    private void startTimeoutChecker() {
        final long interval = Math.max(checkInterval, 3000L); // minimum 3 seconds
        Runnable timeoutChecker = new Runnable() {
            @Override
            public void run() {
                for (sleepInterval(interval); !shutdown; sleepInterval(interval)) {
                    long current = System.currentTimeMillis();
                    for (HiveSession session : new ArrayList<HiveSession>(handleToSession.values())) {
                        if (sessionTimeout > 0 && session.getLastAccessTime() + sessionTimeout <= current
                                && (!checkOperation || session.getNoOperationTime() > sessionTimeout)) {
                            SessionHandle handle = session.getSessionHandle();
                            LOG.warn("Session " + handle + " is Timed-out (last access : "
                                    + new Date(session.getLastAccessTime()) + ") and will be closed");
                            try {
                                closeSession(handle);
                            } catch (HiveSQLException e) {
                                LOG.warn("Exception is thrown closing session " + handle, e);
                            }
                        } else {
                            session.closeExpiredOperations();
                        }
                    }
                }
            }

            private void sleepInterval(long interval) {
                try {
                    Thread.sleep(interval);
                } catch (InterruptedException e) {
                    // ignore
                }
            }
        };
        backgroundOperationPool.execute(timeoutChecker);
    }

    @Override
    public synchronized void stop() {
        super.stop();
        shutdown = true;
        if (backgroundOperationPool != null) {
            backgroundOperationPool.shutdown();
            long timeout = hiveConf.getTimeVar(ConfVars.HIVE_SERVER2_ASYNC_EXEC_SHUTDOWN_TIMEOUT, TimeUnit.SECONDS);
            try {
                backgroundOperationPool.awaitTermination(timeout, TimeUnit.SECONDS);
            } catch (InterruptedException e) {
                LOG.warn("HIVE_SERVER2_ASYNC_EXEC_SHUTDOWN_TIMEOUT = " + timeout
                        + " seconds has been exceeded. RUNNING background operations will be shut down", e);
            }
            backgroundOperationPool = null;
        }
        cleanupLoggingRootDir();
    }

    private void cleanupLoggingRootDir() {
        if (isOperationLogEnabled) {
            try {
                FileUtils.forceDelete(operationLogRootDir);
            } catch (Exception e) {
                LOG.warn("Failed to cleanup root dir of HS2 logging: " + operationLogRootDir.getAbsolutePath(), e);
            }
        }
    }

    public SessionHandle openSession(TProtocolVersion protocol, String username, String password, String ipAddress,
            Map<String, String> sessionConf) throws HiveSQLException {
        return openSession(protocol, username, password, ipAddress, sessionConf, false, null);
    }

    /**
     * Opens a new session and creates a session handle.
     * The username passed to this method is the effective username.
     * If withImpersonation is true (==doAs true) we wrap all the calls in HiveSession
     * within a UGI.doAs, where UGI corresponds to the effective user.
     * @see org.apache.hive.service.cli.thrift.ThriftCLIService#getUserName()
     *
     * @param protocol
     * @param username
     * @param password
     * @param ipAddress
     * @param sessionConf
     * @param withImpersonation
     * @param delegationToken
     * @return
     * @throws HiveSQLException
     */
    public SessionHandle openSession(TProtocolVersion protocol, String username, String password, String ipAddress,
            Map<String, String> sessionConf, boolean withImpersonation, String delegationToken)
            throws HiveSQLException {
        HiveSession session;
        // If doAs is set to true for HiveServer2, we will create a proxy object for the session impl.
        // Within the proxy object, we wrap the method call in a UserGroupInformation#doAs
        if (withImpersonation) {
            HiveSessionImplwithUGI sessionWithUGI = new HiveSessionImplwithUGI(protocol, username, password,
                    hiveConf, ipAddress, delegationToken);
            session = HiveSessionProxy.getProxy(sessionWithUGI, sessionWithUGI.getSessionUgi());
            sessionWithUGI.setProxySession(session);
        } else {
            session = new HiveSessionImpl(protocol, username, password, hiveConf, ipAddress);
        }
        session.setSessionManager(this);
        session.setOperationManager(operationManager);
        try {
            session.open(sessionConf);
        } catch (Exception e) {
            try {
                session.close();
            } catch (Throwable t) {
                LOG.warn("Error closing session", t);
            }
            session = null;
            throw new HiveSQLException("Failed to open new session: " + e, e);
        }
        if (isOperationLogEnabled) {
            session.setOperationLogSessionDir(operationLogRootDir);
        }
        handleToSession.put(session.getSessionHandle(), session);
        return session.getSessionHandle();
    }

    public void closeSession(SessionHandle sessionHandle) throws HiveSQLException {
        HiveSession session = handleToSession.remove(sessionHandle);
        if (session == null) {
            throw new HiveSQLException("Session does not exist!");
        }
        session.close();
    }

    public HiveSession getSession(SessionHandle sessionHandle) throws HiveSQLException {
        HiveSession session = handleToSession.get(sessionHandle);
        if (session == null) {
            throw new HiveSQLException("Invalid SessionHandle: " + sessionHandle);
        }
        return session;
    }

    public OperationManager getOperationManager() {
        return operationManager;
    }

    private static ThreadLocal<String> threadLocalIpAddress = new ThreadLocal<String>() {
        @Override
        protected synchronized String initialValue() {
            return null;
        }
    };

    public static void setIpAddress(String ipAddress) {
        threadLocalIpAddress.set(ipAddress);
    }

    public static void clearIpAddress() {
        threadLocalIpAddress.remove();
    }

    public static String getIpAddress() {
        return threadLocalIpAddress.get();
    }

    private static ThreadLocal<String> threadLocalUserName = new ThreadLocal<String>() {
        @Override
        protected synchronized String initialValue() {
            return null;
        }
    };

    public static void setUserName(String userName) {
        threadLocalUserName.set(userName);
    }

    public static void clearUserName() {
        threadLocalUserName.remove();
    }

    public static String getUserName() {
        return threadLocalUserName.get();
    }

    private static ThreadLocal<String> threadLocalProxyUserName = new ThreadLocal<String>() {
        @Override
        protected synchronized String initialValue() {
            return null;
        }
    };

    public static void setProxyUserName(String userName) {
        LOG.debug("setting proxy user name based on query param to: " + userName);
        threadLocalProxyUserName.set(userName);
    }

    public static String getProxyUserName() {
        return threadLocalProxyUserName.get();
    }

    public static void clearProxyUserName() {
        threadLocalProxyUserName.remove();
    }

    public Future<?> submitBackgroundOperation(Runnable r) {
        return backgroundOperationPool.submit(r);
    }

    public int getOpenSessionCount() {
        return handleToSession.size();
    }
}