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.base;
  
  import java.util.concurrent.BlockingQueue;
  import java.util.concurrent.ExecutionException;
  import java.util.concurrent.Future;
  import java.util.concurrent.ScheduledThreadPoolExecutor;
  import java.util.concurrent.ThreadPoolExecutor;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.atomic.AtomicLong;
  
  import com.google.common.util.concurrent.ThreadFactoryBuilder;
  
  import org.apache.aurora.common.application.Lifecycle;
  import org.apache.aurora.common.stats.Stats;
  import org.slf4j.Logger;
  
  import static java.util.Objects.requireNonNull;
  
  /**
   * Utility class for facilitating async scheduling.
   */
  public final class AsyncUtil {
  
    private AsyncUtil() {
      // Utility class.
    }
  
    private static final AtomicLong UNCAUGHT_EXCEPTIONS =
        Stats.exportLong("async_executor_uncaught_exceptions");
  
    /**
     * Creates a {@link ScheduledThreadPoolExecutor} that logs unhandled errors.
     *
     * @param poolSize Thread pool size.
     * @param nameFormat Thread naming format.
     * @param logger Logger instance.
     * @return instance of {@link ScheduledThreadPoolExecutor} enabled to log unhandled exceptions.
     */
    public static ScheduledThreadPoolExecutor loggingScheduledExecutor(
        int poolSize,
        String nameFormat,
        final Logger logger) {
  
      requireNonNull(nameFormat);
  
      return new ScheduledThreadPoolExecutor(
          poolSize,
          new ThreadFactoryBuilder().setDaemon(true).setNameFormat(nameFormat).build()) {
            @Override
            protected void afterExecute(Runnable runnable, Throwable throwable) {
              super.afterExecute(runnable, throwable);
              evaluateResult(runnable, throwable, logger);
            }
          };
    }
  
    /**
     * Creates a single-threaded {@link ScheduledThreadPoolExecutor} that logs unhandled errors.
     *
     * @param nameFormat Thread naming format.
     * @param logger Logger instance.
     * @return instance of {@link ScheduledThreadPoolExecutor} enabled to log unhandled exceptions.
     */
    public static ScheduledThreadPoolExecutor singleThreadLoggingScheduledExecutor(
        String nameFormat,
        Logger logger) {
  
      return loggingScheduledExecutor(1, nameFormat, logger);
    }
  
    /**
     * Creates a {@link ThreadPoolExecutor} that logs unhandled errors.
     *
     * @param corePoolSize see {@link ThreadPoolExecutor}.
     * @param maxPoolSize see {@link ThreadPoolExecutor}.
     * @param workQueue see {@link ThreadPoolExecutor}.
     * @param nameFormat Thread naming format.
     * @param logger Logger instance.
     * @return instance of {@link ThreadPoolExecutor} enabled to log unhandled exceptions.
     */
    public static ThreadPoolExecutor loggingExecutor(
        int corePoolSize,
        int maxPoolSize,
        BlockingQueue<Runnable> workQueue,
        String nameFormat,
        final Logger logger) {
  
      return new ThreadPoolExecutor(
          corePoolSize,
          maxPoolSize,
          0L,
          TimeUnit.MILLISECONDS,
          workQueue,
          new ThreadFactoryBuilder().setDaemon(true).setNameFormat(nameFormat).build()) {
            @Override
            protected void afterExecute(Runnable runnable, Throwable throwable) {
              super.afterExecute(runnable, throwable);
              evaluateResult(runnable, throwable, logger);
            }
          };
    }
  
    /**
     * Helper wrapper to call the provided {@link Lifecycle} on unhandled error.
     *
     * @param lifecycle {@link Lifecycle} instance.
     * @param logger Logger instance.
     * @param message message to log.
     * @param runnable {@link Runnable} to wrap.
     * @return A new {@link Runnable} logging an error and calling {@link Lifecycle#shutdown()}.
     */
    public static Runnable shutdownOnError(
        Lifecycle lifecycle,
        Logger logger,
        String message,
        Runnable runnable) {
  
      requireNonNull(lifecycle);
      requireNonNull(logger);
      requireNonNull(message);
      requireNonNull(runnable);
  
      return () -> {
        try {
          runnable.run();
        } catch (Throwable t) {
          logger.error(message, t);
          lifecycle.shutdown();
        }
      };
    }
  
    private static void evaluateResult(Runnable runnable, Throwable throwable, Logger logger) {
      // See java.util.concurrent.ThreadPoolExecutor#afterExecute(Runnable, Throwable)
      // for more details and an implementation example.
      if (throwable == null) {
        if (runnable instanceof Future) {
          try {
            Future<?> future = (Future<?>) runnable;
            if (future.isDone()) {
              future.get();
            }
          } catch (InterruptedException ie) {
            Thread.currentThread().interrupt();
          } catch (ExecutionException ee) {
            logger.error(ee.toString(), ee);
            UNCAUGHT_EXCEPTIONS.incrementAndGet();
          }
        }
      } else {
        logger.error(throwable.toString(), throwable);
        UNCAUGHT_EXCEPTIONS.incrementAndGet();
      }
    }
  }
  