alluxio-master-2.x/table/server/underdb/hive/src/main/java/alluxio/table/under/hive/util/DefaultHiveClientPool.java
/*
 * The Alluxio Open Foundation licenses this work under the Apache License, version 2.0
 * (the "License"). You may not use this work except in compliance with the License, which is
 * available at www.apache.org/licenses/LICENSE-2.0
 *
 * This software is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
 * either express or implied, as more fully set forth in the License.
 *
 * See the NOTICE file distributed with this work for information regarding copyright ownership.
 */

package alluxio.table.under.hive.util;

import static alluxio.conf.PropertyKey.TABLE_UDB_HIVE_CLIENTPOOL_MAX;
import static alluxio.conf.PropertyKey.TABLE_UDB_HIVE_CLIENTPOOL_MIN;

import alluxio.Constants;
import alluxio.conf.Configuration;
import alluxio.metrics.MetricKey;
import alluxio.metrics.MetricsSystem;
import alluxio.resource.CloseableResource;
import alluxio.util.ThreadFactoryUtils;

import com.codahale.metrics.Counter;
import org.apache.hadoop.hive.conf.HiveConf;
import org.apache.hadoop.hive.metastore.HiveMetaHookLoader;
import org.apache.hadoop.hive.metastore.HiveMetaStoreClient;
import org.apache.hadoop.hive.metastore.IMetaStoreClient;
import org.apache.hadoop.hive.metastore.RetryingMetaStoreClient;
import org.apache.thrift.TException;

import java.io.IOException;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import javax.annotation.concurrent.ThreadSafe;

/**
 * A pool for hive clients, since hive clients are not thread safe.
 */
@ThreadSafe
public final class DefaultHiveClientPool extends AbstractHiveClientPool {
  private static final ScheduledExecutorService GC_EXECUTOR =
      new ScheduledThreadPoolExecutor(1, ThreadFactoryUtils.build("HiveClientPool-GC-%d", true));
  private static final HiveMetaHookLoader NOOP_HOOK = table -> null;
  private static final Counter COUNTER = MetricsSystem.counter(
      MetricKey.CLIENT_DEFAULT_HIVE_CLIENT_COUNT.getName());

  private final long mGcThresholdMs;
  private final String mConnectionUri;

  /**
   * Creates a new hive client client pool.
   *
   * @param connectionUri the connect uri for the hive metastore
   */
  public DefaultHiveClientPool(String connectionUri) {
    super(Options.defaultOptions()
        .setMinCapacity(Configuration.getInt(TABLE_UDB_HIVE_CLIENTPOOL_MIN))
        .setMaxCapacity(Configuration.getInt(TABLE_UDB_HIVE_CLIENTPOOL_MAX))
        .setGcIntervalMs(5L * Constants.MINUTE_MS)
        .setGcExecutor(GC_EXECUTOR));
    mConnectionUri = connectionUri;
    mGcThresholdMs = 5L * Constants.MINUTE_MS;
  }

  @Override
  protected void closeResource(IMetaStoreClient client) {
    client.close();
  }

  @Override
  protected IMetaStoreClient createNewResource() throws IOException {
    // Hive uses/saves the thread context class loader.
    ClassLoader currentClassLoader = Thread.currentThread().getContextClassLoader();
    try {
      // use the extension class loader
      Thread.currentThread().setContextClassLoader(this.getClass().getClassLoader());
      HiveConf conf = new HiveConf();
      conf.verifyAndSet("hive.metastore.uris", mConnectionUri);

      IMetaStoreClient client = HMSClientFactory.newInstance(
          RetryingMetaStoreClient.getProxy(conf, NOOP_HOOK, HiveMetaStoreClient.class.getName()));
      return client;
    } catch (NullPointerException | TException e) {
      // HiveMetaStoreClient throws a NPE if the uri is not a uri for hive metastore
      throw new IOException(String
          .format("Failed to create client to hive metastore: %s. error: %s", mConnectionUri,
              e.getMessage()), e);
    } finally {
      Thread.currentThread().setContextClassLoader(currentClassLoader);
    }
  }

  @Override
  protected boolean isHealthy(IMetaStoreClient client) {
    // there is no way to check if a hive client is connected.
    // TODO(gpang): periodically and asynchronously check the health of clients
    return true;
  }

  @Override
  protected Counter getMetricCounter() {
    return COUNTER;
  }

  @Override
  protected boolean shouldGc(ResourceInternal<IMetaStoreClient> clientResourceInternal) {
    return System.currentTimeMillis() - clientResourceInternal
        .getLastAccessTimeMs() > mGcThresholdMs;
  }

  @Override
  public CloseableResource<IMetaStoreClient> acquireClientResource() throws IOException {
    return new CloseableResource<IMetaStoreClient>(acquire()) {
      @Override
      public void closeResource() {
        release(get());
      }
    };
  }
}
