src/main/java/org/apache/aurora/GuiceUtils.java

/**
 * 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 org.apache.aurora;

import java.lang.annotation.Retention;
import java.lang.annotation.Target;
import java.lang.reflect.Method;
import java.util.Set;

import javax.inject.Qualifier;

import com.google.common.base.Function;
import com.google.common.base.Preconditions;
import com.google.common.base.Predicate;
import com.google.common.base.Predicates;
import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
import com.google.common.cache.LoadingCache;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Iterables;
import com.google.inject.Binder;
import com.google.inject.matcher.AbstractMatcher;
import com.google.inject.matcher.Matcher;
import com.google.inject.matcher.Matchers;

import org.aopalliance.intercept.MethodInterceptor;
import org.aopalliance.intercept.MethodInvocation;
import org.apache.aurora.common.collections.Pair;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import static java.lang.annotation.ElementType.METHOD;
import static java.lang.annotation.RetentionPolicy.RUNTIME;

/**
 * Utilities for guice configuration in aurora.
 */
public final class GuiceUtils {

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

  // Method annotation that allows a trapped interface to whitelist methods that may throw
  // unchecked exceptions.
  @Qualifier
  @Target(METHOD) @Retention(RUNTIME)
  public @interface AllowUnchecked { }

  private GuiceUtils() {
    // utility
  }

  private static final Function<Method, Pair<String, Class<?>[]>> CANONICALIZE =
      method -> Pair.of(method.getName(), method.getParameterTypes());

  /**
   * Creates a matcher that will match methods of an interface, optionally excluding inherited
   * methods.
   *
   * @param matchInterface The interface to match.
   * @param declaredMethodsOnly if {@code true} only methods directly declared in the interface
   *                            will be matched, otherwise all methods on the interface are matched.
   * @return A new matcher instance.
   */
  public static Matcher<Method> interfaceMatcher(
      Class<?> matchInterface,
      boolean declaredMethodsOnly) {

    Method[] methods =
        declaredMethodsOnly ? matchInterface.getDeclaredMethods() : matchInterface.getMethods();
    final Set<Pair<String, Class<?>[]>> interfaceMethods =
        ImmutableSet.copyOf(Iterables.transform(ImmutableList.copyOf(methods), CANONICALIZE));
    final LoadingCache<Method, Pair<String, Class<?>[]>> cache = CacheBuilder.newBuilder()
        .build(CacheLoader.from(CANONICALIZE));

    return new AbstractMatcher<Method>() {
      @Override
      public boolean matches(Method method) {
        return interfaceMethods.contains(cache.getUnchecked(method));
      }
    };
  }

  /**
   * Binds an interceptor that ensures the main ClassLoader is bound as the thread context
   * {@link ClassLoader} during JNI callbacks from mesos.  Some libraries require a thread
   * context ClassLoader be set and this ensures those libraries work properly.
   *
   * @param binder The binder to use to register an interceptor with.
   * @param wrapInterface Interface whose methods should wrapped.
   */
  public static void bindJNIContextClassLoader(Binder binder, Class<?> wrapInterface) {
    final ClassLoader mainClassLoader = GuiceUtils.class.getClassLoader();
    binder.bindInterceptor(
        Matchers.subclassesOf(wrapInterface),
        interfaceMatcher(wrapInterface, false),
        new MethodInterceptor() {
          @Override
          public Object invoke(MethodInvocation invocation) throws Throwable {
            Thread currentThread = Thread.currentThread();
            ClassLoader prior = currentThread.getContextClassLoader();
            try {
              currentThread.setContextClassLoader(mainClassLoader);
              return invocation.proceed();
            } finally {
              currentThread.setContextClassLoader(prior);
            }
          }
        });
  }

  private static final Predicate<Method> IS_WHITELISTED =
      method -> method.getAnnotation(AllowUnchecked.class) != null;

  private static final Matcher<Method> WHITELIST_MATCHER = new AbstractMatcher<Method>() {
    @Override
    public boolean matches(Method method) {
      return IS_WHITELISTED.apply(method);
    }
  };

  private static final Predicate<Method> VOID_METHOD =
      method -> method.getReturnType() == Void.TYPE;

  /**
   * Binds an exception trap on all interface methods of all classes bound against an interface.
   * Individual methods may opt out of trapping by annotating with {@link AllowUnchecked}.
   * Only void methods are allowed, any non-void interface methods must explicitly opt out.
   *
   * @param binder The binder to register an interceptor with.
   * @param wrapInterface Interface whose methods should be wrapped.
   * @throws IllegalArgumentException If any of the non-whitelisted interface methods are non-void.
   */
  public static void bindExceptionTrap(Binder binder, Class<?> wrapInterface)
      throws IllegalArgumentException {

    Set<Method> disallowed = ImmutableSet.copyOf(Iterables.filter(
        ImmutableList.copyOf(wrapInterface.getMethods()),
        Predicates.and(Predicates.not(IS_WHITELISTED), Predicates.not(VOID_METHOD))));
    Preconditions.checkArgument(disallowed.isEmpty(),
        "Non-void methods must be explicitly whitelisted with @AllowUnchecked: %s", disallowed);

    Matcher<Method> matcher =
        Matchers.not(WHITELIST_MATCHER).and(interfaceMatcher(wrapInterface, false));
    binder.bindInterceptor(Matchers.subclassesOf(wrapInterface), matcher,
        new MethodInterceptor() {
          @Override
          public Object invoke(MethodInvocation invocation) throws Throwable {
            try {
              return invocation.proceed();
            } catch (RuntimeException e) {
              LOG.warn("Trapped uncaught exception: " + e, e);
              return null;
            }
          }
        });
  }
}