content: /**
   * 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.scheduler.thrift.aop;
  
  import java.lang.reflect.Method;
  import java.util.concurrent.atomic.AtomicLong;
  
  import com.google.common.annotations.VisibleForTesting;
  import com.google.common.cache.CacheBuilder;
  import com.google.common.cache.CacheLoader;
  import com.google.common.cache.LoadingCache;
  
  import org.aopalliance.intercept.MethodInterceptor;
  import org.aopalliance.intercept.MethodInvocation;
  import org.apache.aurora.common.stats.SlidingStats;
  import org.apache.aurora.common.stats.Stats;
  import org.apache.aurora.gen.Response;
  import org.apache.aurora.gen.ResponseCode;
  import org.apache.aurora.scheduler.thrift.aop.ThriftWorkload.ThriftWorkloadCounter;
  
  /**
   * A method interceptor that exports counterStats about thrift calls.
   */
  class ThriftStatsExporterInterceptor implements MethodInterceptor {
  
    @VisibleForTesting
    static final String TIMING_STATS_NAME_TEMPLATE = "scheduler_thrift_%s";
    @VisibleForTesting
    static final String WORKLOAD_STATS_NAME_TEMPLATE = "scheduler_thrift_workload_%s";
  
    private final LoadingCache<Method, SlidingStats> timingStats =
        CacheBuilder.newBuilder().build(new CacheLoader<Method, SlidingStats>() {
          @Override
          public SlidingStats load(Method method) {
            return new SlidingStats(
                Stats.normalizeName(String.format(TIMING_STATS_NAME_TEMPLATE, method.getName())),
                "nanos");
          }
        });
  
    private final LoadingCache<Method, AtomicLong> workloadStats =
        CacheBuilder.newBuilder().build(new CacheLoader<Method, AtomicLong>() {
          @Override
          public AtomicLong load(Method method) {
            return Stats.exportLong(
                Stats.normalizeName(String.format(WORKLOAD_STATS_NAME_TEMPLATE, method.getName())));
          }
        });
  
    @Override
    public Object invoke(MethodInvocation invocation) throws Throwable {
      Method method = invocation.getMethod();
      SlidingStats stat = timingStats.getUnchecked(method);
      long start = System.nanoTime();
      Response response = null;
      try {
        response = (Response) invocation.proceed();
      } finally {
        stat.accumulate(System.nanoTime() - start);
        if (response != null
            && response.getResponseCode() == ResponseCode.OK
            && method.isAnnotationPresent(ThriftWorkload.class)) {
  
          ThriftWorkloadCounter counter = method.getAnnotation(ThriftWorkload.class)
              .value()
              .newInstance();
          workloadStats.getUnchecked(method).addAndGet(counter.apply(response.getResult()));
        }
      }
      return response;
    }
  }
  