/*
 * JORAM: Java(TM) Open Reliable Asynchronous Messaging
 * Copyright (C) 2019 ScalAgent Distributed Technologies
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or any later version.
 * 
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 * 
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307
 * USA.
 *
 * Initial developer(s): ScalAgent Distributed Technologies
 * Contributor(s): 
 */
package fr.dyade.aaa.common.concurrent;

import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

/**
 * Imposes a time limit on method calls.
 */
public class TimeLimiter {
  /**
   * Invokes a specified Callable with the specified time limit.
   * If the target method call finishes before the limit is reached, the return value or the exception is
   * propagated. If the time limit is reached, we attempt to abort the call, and throw a {@link TimeoutException}
   * to the caller.
   *
   * @param <T> the type of the task's result
   * @param callable the Callable to execute
   * @param timeoutDuration with timeoutUnit, the maximum length of time to wait
   * @param timeoutUnit with timeoutDuration, the maximum length of time to wait
   * @return the result returned by the Callable
   * @throws Exception if the time limit is reached, the thread is interrupted, or an exception
   * is thrown during execution.
   */
  public static <T> T callWithTimeout(Callable<T> callable, long timeoutDuration, TimeUnit timeoutUnit) throws Exception {
    final ExecutorService executor = Executors.newSingleThreadExecutor();
    Future<T> future = executor.submit(callable);
    executor.shutdown(); // This does not cancel the already-scheduled task.

    try {
      return future.get(timeoutDuration, timeoutUnit);
    } catch (InterruptedException | TimeoutException e) {
      future.cancel(true /* mayInterruptIfRunning */);
      throw e;
    } catch (ExecutionException e) {
      //unwrap the root cause
      Throwable t = e.getCause();
      if (t instanceof Error) {
        throw (Error) t;
      } else if (t instanceof Exception) {
        throw (Exception) t;
      } else {
        throw new IllegalStateException(t);
      }
    }
  }

  /**
   * Invokes a specified Runnable with the specified time limit.
   * If the target method call finishes before the limit is reached, the return value or the exception is
   * propagated. If the time limit is reached, we attempt to abort the call, and throw a {@link TimeoutException}
   * to the caller.
   *
   * @param runnable the Runnable to execute
   * @param timeoutDuration with timeoutUnit, the maximum length of time to wait
   * @param timeoutUnit with timeoutDuration, the maximum length of time to wait
   * @throws Exception if the time limit is reached, the thread is interrupted, or an exception
   * is thrown during execution.
   */
  public void runWithTimeout(Runnable runnable, long timeoutDuration, TimeUnit timeoutUnit) throws Exception {
    final ExecutorService executor = Executors.newSingleThreadExecutor();
    Future<?> future = executor.submit(runnable);
    executor.shutdown(); // This does not cancel the already-scheduled task.

    try {
      future.get(timeoutDuration, timeoutUnit);
    } catch (InterruptedException | TimeoutException e) {
      future.cancel(true /* mayInterruptIfRunning */);
      throw e;
    } catch (ExecutionException e) {
      //unwrap the root cause
      Throwable t = e.getCause();
      if (t instanceof Error) {
        throw (Error) t;
      } else if (t instanceof Exception) {
        throw (Exception) t;
      } else {
        throw new IllegalStateException(t);
      }
    }
  }

}
