RequestPoolService.java

Go to the documentation of this file.

// Copyright 2014 Cloudera Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package com.cloudera.impala.util;

import java.io.File;
import java.io.IOException;
import java.net.MalformedURLException;
import java.net.URL;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.security.UserGroupInformation;
import org.apache.hadoop.yarn.api.records.QueueACL;
import org.apache.hadoop.yarn.conf.YarnConfiguration;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.fair.AllocationConfiguration;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.fair.AllocationFileLoaderService;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.fair.FairSchedulerConfiguration;
import org.apache.thrift.TException;
import org.apache.thrift.TSerializer;
import org.apache.thrift.protocol.TBinaryProtocol;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.cloudera.impala.authorization.User;
import com.cloudera.impala.common.ByteUnits;
import com.cloudera.impala.common.ImpalaException;
import com.cloudera.impala.common.InternalException;
import com.cloudera.impala.common.JniUtil;
import com.cloudera.impala.thrift.TErrorCode;
import com.cloudera.impala.thrift.TPoolConfigParams;
import com.cloudera.impala.thrift.TPoolConfigResult;
import com.cloudera.impala.thrift.TResolveRequestPoolParams;
import com.cloudera.impala.thrift.TResolveRequestPoolResult;
import com.cloudera.impala.thrift.TStatus;
import com.cloudera.impala.util.FileWatchService.FileChangeListener;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import com.google.common.base.Strings;
import com.google.common.collect.Lists;

public class RequestPoolService {
  final static Logger LOG = LoggerFactory.getLogger(RequestPoolService.class);

  private final static TBinaryProtocol.Factory protocolFactory_ =
      new TBinaryProtocol.Factory();
  // Used to ensure start() has been called before any other methods can be used.
  private final AtomicBoolean running_;

  // Key for the default maximum number of running queries ("placed reservations")
  // property. The per-pool key name is this key with the pool name appended, e.g.
  // "{key}.{pool}".
  final static String LLAMA_MAX_PLACED_RESERVATIONS_KEY =
      "llama.am.throttling.maximum.placed.reservations";

  // Default value for the maximum.placed.reservations property. Note that this value
  // differs from the current Llama default of 10000 which is too high.
  final static int LLAMA_MAX_PLACED_RESERVATIONS_DEFAULT = 200;

  // Key for the default maximum number of queued requests ("queued reservations")
  // property. The per-pool key name is this key with the pool name appended, e.g.
  // "{key}.{pool}".
  final static String LLAMA_MAX_QUEUED_RESERVATIONS_KEY =
      "llama.am.throttling.maximum.queued.reservations";

  // Default value for the maximum.queued.reservations property. Note that this value
  // differs from the current Llama default of 0 which disables queuing.
  final static int LLAMA_MAX_QUEUED_RESERVATIONS_DEFAULT = 200;

  // String format for a per-pool configuration key. First parameter is the key for the
  // default, e.g. LLAMA_MAX_PLACED_RESERVATIONS_KEY, and the second parameter is the
  // pool name.
  final static String LLAMA_PER_POOL_CONFIG_KEY_FORMAT = "%s.%s";

  // Watches for changes to the fair scheduler allocation file.
  @VisibleForTesting
  final AllocationFileLoaderService allocLoader_;

  // Provides access to the fair scheduler allocation file. An AtomicReference becaus it
  // is reset when the allocation configuration file changes and other threads access it.
  private final AtomicReference<AllocationConfiguration> allocationConf_;

  // Watches the Llama configuration file for changes.
  @VisibleForTesting
  final FileWatchService llamaConfWatcher_;

  // Used by this class to access to the configs provided by the Llama configuration.
  // This is replaced when the Llama configuration file changes.
  private volatile Configuration llamaConf_;

  // URL of the Llama configuration file.
  private final URL llamaConfUrl_;

  private final class LlamaConfWatcher implements FileChangeListener {
    public void onFileChange() {
      // If llamaConfUrl_ is null the watcher should not have been created.
      Preconditions.checkNotNull(llamaConfUrl_);
      LOG.info("Loading Llama configuration: " + llamaConfUrl_.getFile());
      Configuration conf = new Configuration();
      conf.addResource(llamaConfUrl_);
      llamaConf_ = conf;
    }
  }

  public RequestPoolService(final String fsAllocationPath, final String llamaSitePath) {
    Preconditions.checkNotNull(fsAllocationPath);
    running_ = new AtomicBoolean(false);
    allocationConf_ = new AtomicReference<AllocationConfiguration>();
    URL fsAllocationURL = getURL(fsAllocationPath);
    if (fsAllocationURL == null) {
      throw new IllegalArgumentException(
          "Unable to find allocation configuration file: " + fsAllocationPath);
    }
    Configuration allocConf = new Configuration(false);
    allocConf.set(FairSchedulerConfiguration.ALLOCATION_FILE, fsAllocationURL.getPath());
    allocLoader_ = new AllocationFileLoaderService();
    allocLoader_.init(allocConf);

    if (!Strings.isNullOrEmpty(llamaSitePath)) {
      llamaConfUrl_ = getURL(llamaSitePath);
      if (llamaConfUrl_ == null) {
        throw new IllegalArgumentException(
            "Unable to find Llama configuration file: " + llamaSitePath);
      }
      llamaConf_ = new Configuration(false);
      llamaConf_.addResource(llamaConfUrl_);
      llamaConfWatcher_ = new FileWatchService(new File(llamaConfUrl_.getPath()),
          new LlamaConfWatcher());
    } else {
      llamaConfWatcher_ = null;
      llamaConfUrl_ = null;
    }
  }

  @VisibleForTesting
  static URL getURL(String path) {
    Preconditions.checkNotNull(path);
    File file = new File(path);
    file = file.getAbsoluteFile();
    if (!file.exists()) {
      LOG.error("Unable to find specified file: " + path);
      return null;
    }
    try {
      return file.toURI().toURL();
    } catch (MalformedURLException ex) {
      LOG.error("Unable to construct URL for file: " + path, ex);
      return null;
    }
  }

  public void start() {
    Preconditions.checkState(!running_.get());
    allocLoader_.setReloadListener(new AllocationFileLoaderService.Listener() {
      @Override
      public void onReload(AllocationConfiguration info) {
        allocationConf_.set(info);
      }
    });
    allocLoader_.start();
    try {
      allocLoader_.reloadAllocations();
    } catch (Exception ex) {
      try {
        stopInternal();
      } catch (Exception stopEx) {
        LOG.error("Unable to stop AllocationFileLoaderService after failed start.",
            stopEx);
      }
      throw new RuntimeException(ex);
    }
    if (llamaConfWatcher_ != null) llamaConfWatcher_.start();
    running_.set(true);
  }

  public void stop() {
    Preconditions.checkState(running_.get());
    stopInternal();
  }

  private void stopInternal() {
    running_.set(false);
    if (llamaConfWatcher_ != null) llamaConfWatcher_.stop();
    allocLoader_.stop();
  }

  public byte[] resolveRequestPool(byte[] thriftResolvePoolParams)
      throws ImpalaException {
    TResolveRequestPoolParams resolvePoolParams = new TResolveRequestPoolParams();
    JniUtil.deserializeThrift(protocolFactory_, resolvePoolParams,
        thriftResolvePoolParams);
    TResolveRequestPoolResult result = resolveRequestPool(resolvePoolParams);
    LOG.trace("resolveRequestPool(pool={}, user={}): resolved_pool={}, has_access={}",
        new Object[] { resolvePoolParams.getRequested_pool(), resolvePoolParams.getUser(),
                       result.resolved_pool, result.has_access });
    try {
      return new TSerializer(protocolFactory_).serialize(result);
    } catch (TException e) {
      throw new InternalException(e.getMessage());
    }
  }

  @VisibleForTesting
  TResolveRequestPoolResult resolveRequestPool(
      TResolveRequestPoolParams resolvePoolParams) {
    String requestedPool = resolvePoolParams.getRequested_pool();
    String user = resolvePoolParams.getUser();
    TResolveRequestPoolResult result = new TResolveRequestPoolResult();
    String errorMessage = null;
    String pool = null;
    try {
      pool = assignToPool(requestedPool, user);
    } catch (IOException ex) {
      errorMessage = ex.getMessage();
      if (errorMessage.startsWith("No groups found for user")) {
        // The error thrown when using the 'primaryGroup' or 'secondaryGroup' rules and
        // the user does not exist are not helpful.
        errorMessage = String.format(
            "Failed to resolve user '%s' to a pool while evaluating the " +
            "'primaryGroup' or 'secondaryGroup' queue placement rules because no " +
            "groups were found for the user. This is likely because the user does not " +
            "exist on the local operating system.", resolvePoolParams.getUser());
      }
      LOG.warn(String.format("Error assigning to pool. requested='%s', user='%s', msg=%s",
          requestedPool, user, errorMessage), ex);
    }
    if (pool == null) {
      if (errorMessage == null) {
        // This occurs when assignToPool returns null (not an error), i.e. if the pool
        // cannot be resolved according to the policy.
        result.setStatus(new TStatus(TErrorCode.OK, Lists.<String>newArrayList()));
      } else {
        // If Yarn throws an exception, return an error status.
        result.setStatus(
            new TStatus(TErrorCode.INTERNAL_ERROR, Lists.newArrayList(errorMessage)));
      }
    } else {
      result.setResolved_pool(pool);
      result.setHas_access(hasAccess(pool, user));
      result.setStatus(new TStatus(TErrorCode.OK, Lists.<String>newArrayList()));
    }
    return result;
  }

  public byte[] getPoolConfig(byte[] thriftPoolConfigParams) throws ImpalaException {
    Preconditions.checkState(running_.get());
    TPoolConfigParams poolConfigParams = new TPoolConfigParams();
    JniUtil.deserializeThrift(protocolFactory_, poolConfigParams,
        thriftPoolConfigParams);
    TPoolConfigResult result = getPoolConfig(poolConfigParams.getPool());
    try {
      return new TSerializer(protocolFactory_).serialize(result);
    } catch (TException e) {
      throw new InternalException(e.getMessage());
    }
  }

  @VisibleForTesting
  TPoolConfigResult getPoolConfig(String pool) {
    TPoolConfigResult result = new TPoolConfigResult();
    int maxMemoryMb = allocationConf_.get().getMaxResources(pool).getMemory();
    result.setMem_limit(
        maxMemoryMb == Integer.MAX_VALUE ? -1 : (long) maxMemoryMb * ByteUnits.MEGABYTE);
    if (llamaConf_ == null) {
      result.setMax_requests(LLAMA_MAX_PLACED_RESERVATIONS_DEFAULT);
      result.setMax_queued(LLAMA_MAX_QUEUED_RESERVATIONS_DEFAULT);
    } else {
      // Capture the current llamaConf_ in case it changes while we're using it.
      Configuration currentLlamaConf = llamaConf_;
      result.setMax_requests(getLlamaPoolConfigValue(currentLlamaConf, pool,
          LLAMA_MAX_PLACED_RESERVATIONS_KEY,
          LLAMA_MAX_PLACED_RESERVATIONS_DEFAULT));
      result.setMax_queued(getLlamaPoolConfigValue(currentLlamaConf, pool,
          LLAMA_MAX_QUEUED_RESERVATIONS_KEY,
          LLAMA_MAX_QUEUED_RESERVATIONS_DEFAULT));
    }
    LOG.trace("getPoolConfig(pool={}): mem_limit={}, max_requests={}, max_queued={}",
        new Object[] { pool, result.mem_limit, result.max_requests, result.max_queued });
    return result;
  }

  private int getLlamaPoolConfigValue(Configuration conf, String pool, String key,
      int defaultValue) {
    return conf.getInt(String.format(LLAMA_PER_POOL_CONFIG_KEY_FORMAT, key, pool),
        conf.getInt(key, defaultValue));
  }

  @VisibleForTesting
  String assignToPool(String requestedPool, String user)
      throws IOException {
    Preconditions.checkState(running_.get());
    Preconditions.checkNotNull(requestedPool);
    Preconditions.checkArgument(!Strings.isNullOrEmpty(user));
    // Convert the user name to a short name (e.g. 'user1@domain' to 'user1') because
    // assignAppToQueue() will check group membership which should always be done on
    // the short name of the principal.
    String shortName = new User(user).getShortName();
    return allocationConf_.get().getPlacementPolicy().assignAppToQueue(
        requestedPool.isEmpty() ? YarnConfiguration.DEFAULT_QUEUE_NAME : requestedPool,
        shortName);
  }

  @VisibleForTesting
  boolean hasAccess(String pool, String user) {
    Preconditions.checkState(running_.get());
    Preconditions.checkArgument(!Strings.isNullOrEmpty(pool));
    Preconditions.checkArgument(!Strings.isNullOrEmpty(user));
    // Convert the user name to a short name (e.g. 'user1@domain' to 'user1') because
    // the UserGroupInformation will check group membership which should always be done
    // on the short name of the principal.
    String shortName = new User(user).getShortName();
    UserGroupInformation ugi = UserGroupInformation.createRemoteUser(shortName);
    return allocationConf_.get().hasAccess(pool, QueueACL.SUBMIT_APPLICATIONS, ugi);
  }
}