io.cdap.cdap.app.guice.ImpersonatedTwillController Maven / Gradle / Ivy

Show more of this group Show more artifacts with this name
Show all versions of cdap-app-fabric Show documentation
There is a newer version: 6.10.1
/*
 * Copyright © 2016-2017 Cask Data, 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 io.cdap.cdap.app.guice;

import com.google.common.base.Throwables;
import com.google.common.util.concurrent.Futures;
import io.cdap.cdap.common.ServiceUnavailableException;
import io.cdap.cdap.proto.id.ProgramId;
import io.cdap.cdap.security.impersonation.Impersonator;
import org.apache.twill.api.Command;
import org.apache.twill.api.ResourceReport;
import org.apache.twill.api.RunId;
import org.apache.twill.api.ServiceController;
import org.apache.twill.api.TwillController;
import org.apache.twill.api.logging.LogEntry;
import org.apache.twill.api.logging.LogHandler;
import org.apache.twill.discovery.ServiceDiscovered;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Map;
import java.util.Set;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executor;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import javax.annotation.Nullable;

/**
 * A {@link TwillController} wrapper that performs impersonation on {@link #getResourceReport()}, {@link #terminate()}
 * and {@link #kill()}.
 */
final class ImpersonatedTwillController implements TwillController {

  private static final Logger LOG = LoggerFactory.getLogger(ImpersonatedTwillController.class);

  private final TwillController delegate;
  private final Impersonator impersonator;
  private final ProgramId programId;

  ImpersonatedTwillController(TwillController delegate, Impersonator impersonator, ProgramId programId) {
    this.delegate = delegate;
    this.impersonator = impersonator;
    this.programId = programId;
  }

  @Override
  public void addLogHandler(LogHandler handler) {
    delegate.addLogHandler(handler);
  }

  @Override
  public ServiceDiscovered discoverService(String serviceName) {
    return delegate.discoverService(serviceName);
  }

  @Override
  public Future changeInstances(String runnable, int newCount) {
    return delegate.changeInstances(runnable, newCount);
  }

  @Nullable
  @Override
  public ResourceReport getResourceReport() {
    try {
      return impersonator.doAs(programId, new Callable() {
        @Nullable
        @Override
        public ResourceReport call() throws Exception {
          return delegate.getResourceReport();
        }
      });
    } catch (Exception e) {
      // The delegate.getResourceReport() call never throws exception, hence exception caught must due to
      // impersonation failure.
      if (Throwables.getRootCause(e) instanceof ServiceUnavailableException) {
        // If it is due to some underlying service unavailability, log a debug message and return null
        // It is expected to happen during master process startup
        LOG.debug("Failed in impersonation for program {}", programId, e);
      } else {
        LOG.warn("Unexpected exception in impersonation for program {}", programId, e);
      }
      return null;
    }
  }

  @Override
  public Future restartAllInstances(String runnable) {
    return delegate.restartAllInstances(runnable);
  }

  @Override
  public Future> restartInstances(Map> runnableToInstanceIds) {
    return delegate.restartInstances(runnableToInstanceIds);
  }

  @Override
  public Future restartInstances(String runnable, int instanceId, int... moreInstanceIds) {
    return delegate.restartInstances(runnable, instanceId, moreInstanceIds);
  }

  @Override
  public Future restartInstances(String runnable, Set instanceIds) {
    return delegate.restartInstances(runnable, instanceIds);
  }

  @Override
  public Future> updateLogLevels(Map logLevels) {
    return delegate.updateLogLevels(logLevels);
  }

  @Override
  public Future> updateLogLevels(String runnableName,
                                                             Map logLevels) {
    return delegate.updateLogLevels(runnableName, logLevels);
  }

  @Override
  public Future resetLogLevels(String... loggerNames) {
    return delegate.resetLogLevels(loggerNames);
  }

  @Override
  public Future resetRunnableLogLevels(String runnableName, String... loggerNames) {
    return delegate.resetRunnableLogLevels(runnableName, loggerNames);
  }

  @Override
  public RunId getRunId() {
    return delegate.getRunId();
  }

  @Override
  public Future sendCommand(Command command) {
    return delegate.sendCommand(command);
  }

  @Override
  public Future sendCommand(String runnableName, Command command) {
    return delegate.sendCommand(runnableName, command);
  }

  @Override
  public Future terminate() {
    try {
      return impersonator.doAs(programId, new Callable>() {
        @Override
        public Future call() throws Exception {
          return delegate.terminate();
        }
      });
    } catch (Exception e) {
      return Futures.immediateFailedFuture(e);
    }
  }

  @Override
  public void kill() {
    try {
      impersonator.doAs(programId, new Callable() {
        @Override
        public Void call() throws Exception {
          delegate.kill();
          return null;
        }
      });
    } catch (Exception e) {
      throw Throwables.propagate(e);
    }
  }

  @Override
  public void onRunning(Runnable runnable, Executor executor) {
    delegate.onRunning(runnable, executor);
  }

  @Override
  public void onTerminated(Runnable runnable, Executor executor) {
    delegate.onTerminated(runnable, executor);
  }

  @Override
  public void awaitTerminated() throws ExecutionException {
    delegate.awaitTerminated();
  }

  @Override
  public void awaitTerminated(long timeout, TimeUnit timeoutUnit) throws TimeoutException, ExecutionException {
    delegate.awaitTerminated(timeout, timeoutUnit);
  }

  @Nullable
  @Override
  public TerminationStatus getTerminationStatus() {
    return delegate.getTerminationStatus();
  }
}