gobblin-runtime/src/main/java/org/apache/gobblin/runtime/TaskExecutor.java - gobblin - Git at Google

 /*
  * Licensed to the Apache Software Foundation (ASF) under one or more
  * contributor license agreements.  See the NOTICE file distributed with
  * this work for additional information regarding copyright ownership.
  * The ASF licenses this file to You 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.gobblin.runtime;

 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.Map;
 import java.util.NoSuchElementException;
 import java.util.Properties;
 import java.util.concurrent.ConcurrentSkipListMap;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.SynchronousQueue;
 import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;

 import com.codahale.metrics.Counter;
 import com.codahale.metrics.Gauge;
 import com.codahale.metrics.Meter;
 import com.codahale.metrics.Metric;
 import com.codahale.metrics.MetricSet;
 import com.codahale.metrics.SlidingTimeWindowReservoir;
 import com.codahale.metrics.Timer;
 import com.google.common.base.Optional;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.Maps;
 import com.google.common.util.concurrent.AbstractIdleService;

 import org.apache.hadoop.conf.Configuration;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 import org.apache.gobblin.configuration.ConfigurationKeys;
 import org.apache.gobblin.metrics.GobblinMetrics;
 import org.apache.gobblin.runtime.fork.Fork;
 import org.apache.gobblin.util.ExecutorsUtils;
 import org.apache.gobblin.util.logs.Log4jConfigurationHelper;

 import lombok.Getter;

 import static com.codahale.metrics.MetricRegistry.name;


 /**
  * A class for executing {@link Task}s and retrying failed ones as well as for executing {@link Fork}s.
  *
  * @author Yinan Li
  */
 public class TaskExecutor extends AbstractIdleService {

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

   // Thread pool executor for running tasks
   private final ScheduledExecutorService taskExecutor;

   // A separate thread pool executor for running forks of tasks
   @Getter
   private final ExecutorService forkExecutor;

   // Task retry interval
   @Getter
   private final long retryIntervalInSeconds;

   // The maximum number of items in the queued task time map.
   @Getter
   private final int queuedTaskTimeMaxSize;

   // The maximum age of the items in the queued task time map.
   @Getter
   private final long queuedTaskTimeMaxAge ;

   // Map of queued task ids to queue times.  The key is the task id, the value is the time the task was queued.  If the
   // task is being retried, the time may be in the future.  Entries with time in the future will not be counted as
   // queued until the time is in the past.
   private final Map<String, Long> queuedTasks = Maps.newConcurrentMap();

   // Set of historical queued task times. The key is the UTC epoch time the task started, the value is the milliseconds
   // the task waited to start.
   private final ConcurrentSkipListMap<Long, Long> queuedTaskTimeHistorical = new ConcurrentSkipListMap<>();

   // The timestamp for the last time the metrics were calculated.
   private long lastCalculationTime = 0;

   // The total number of tasks currently queued and queued over the historical lookback period.
   @Getter
   private AtomicInteger queuedTaskCount = new AtomicInteger();

   // The total number of tasks currently queued.
   @Getter
   private AtomicInteger currentQueuedTaskCount = new AtomicInteger();

   // The total number of tasks queued over the historical lookback period.
   @Getter
   private AtomicInteger historicalQueuedTaskCount = new AtomicInteger();

   // The total time tasks have currently been in the queue and were in the queue during the historical lookback period.
   @Getter
   private AtomicLong queuedTaskTotalTime = new AtomicLong();

   // The total time tasks have currently been in the queue.
   @Getter
   private AtomicLong currentQueuedTaskTotalTime = new AtomicLong();

   // The total time tasks have been in the queue during the historical lookback period.
   @Getter
   private AtomicLong historicalQueuedTaskTotalTime = new AtomicLong();

   // Count of running tasks.
   @Getter
   private final Counter runningTaskCount = new Counter();

   // Count of failed tasks.
   @Getter
   private final Meter successfulTaskCount = new Meter();

   // Count of failed tasks.
   @Getter
   private final Meter failedTaskCount = new Meter();

   @Getter
   private final Timer taskCreateAndRunTimer;

   // The metric set exposed from the task executor.
   private final TaskExecutorQueueMetricSet metricSet = new TaskExecutorQueueMetricSet();

   /**
    * Constructor used internally.
    */
   private TaskExecutor(int taskExecutorThreadPoolSize, int coreRetryThreadPoolSize, long retryIntervalInSeconds,
                        int queuedTaskTimeMaxSize, long queuedTaskTimeMaxAge, int timerWindowSize) {
     Preconditions.checkArgument(taskExecutorThreadPoolSize > 0, "Task executor thread pool size should be positive");
     Preconditions.checkArgument(retryIntervalInSeconds > 0, "Task retry interval should be positive");
     Preconditions.checkArgument(queuedTaskTimeMaxSize > 0, "Queued task time max size should be positive");
     Preconditions.checkArgument(queuedTaskTimeMaxAge > 0, "Queued task time max age should be positive");

     // Currently a fixed-size thread pool is used to execute tasks. We probably need to revisit this later.
     this.taskExecutor = ExecutorsUtils.loggingDecorator(Executors.newScheduledThreadPool(
         taskExecutorThreadPoolSize,
         ExecutorsUtils.newThreadFactory(Optional.of(LOG), Optional.of("TaskExecutor-%d"))));

     this.retryIntervalInSeconds = retryIntervalInSeconds;
     this.queuedTaskTimeMaxSize = queuedTaskTimeMaxSize;
     this.queuedTaskTimeMaxAge = queuedTaskTimeMaxAge;
     this.taskCreateAndRunTimer = new Timer(new SlidingTimeWindowReservoir(timerWindowSize, TimeUnit.MINUTES));

     this.forkExecutor = ExecutorsUtils.loggingDecorator(
         new ThreadPoolExecutor(
             // The core thread pool size is equal to that of the task
             // executor as there's at least one fork per task
             taskExecutorThreadPoolSize,
             // The fork executor thread pool size is essentially unbounded. This is to make sure all forks of
             // a task get a thread to run so all forks of the task are making progress. This is necessary since
             // otherwise the parent task will be blocked if the record queue (bounded) of some fork is full and
             // that fork has not yet started to run because of no available thread. The task cannot proceed in
             // this case because it has to make sure every records go to every forks.
             Integer.MAX_VALUE,
             0L,
             TimeUnit.MILLISECONDS,
             // The work queue is a SynchronousQueue. This essentially forces a new thread to be created for each fork.
             new SynchronousQueue<Runnable>(),
             ExecutorsUtils.newThreadFactory(Optional.of(LOG), Optional.of("ForkExecutor-%d"))));
   }

   /**
    * Constructor to work with {@link java.util.Properties}.
    */
   public TaskExecutor(Properties properties) {
     this(Integer.parseInt(properties.getProperty(ConfigurationKeys.TASK_EXECUTOR_THREADPOOL_SIZE_KEY,
             Integer.toString(ConfigurationKeys.DEFAULT_TASK_EXECUTOR_THREADPOOL_SIZE))),
         Integer.parseInt(properties.getProperty(ConfigurationKeys.TASK_RETRY_THREAD_POOL_CORE_SIZE_KEY,
             Integer.toString(ConfigurationKeys.DEFAULT_TASK_RETRY_THREAD_POOL_CORE_SIZE))),
         Long.parseLong(properties.getProperty(ConfigurationKeys.TASK_RETRY_INTERVAL_IN_SEC_KEY,
             Long.toString(ConfigurationKeys.DEFAULT_TASK_RETRY_INTERVAL_IN_SEC))),
         Integer.parseInt(properties.getProperty(ConfigurationKeys.QUEUED_TASK_TIME_MAX_SIZE,
             Integer.toString(ConfigurationKeys.DEFAULT_QUEUED_TASK_TIME_MAX_SIZE))),
         Long.parseLong(properties.getProperty(ConfigurationKeys.QUEUED_TASK_TIME_MAX_AGE,
             Long.toString(ConfigurationKeys.DEFAULT_QUEUED_TASK_TIME_MAX_AGE))),
         Integer.parseInt(properties.getProperty(ConfigurationKeys.METRIC_TIMER_WINDOW_SIZE_IN_MINUTES,
             Integer.toString(ConfigurationKeys.DEFAULT_METRIC_TIMER_WINDOW_SIZE_IN_MINUTES))));
   }

   /**
    * Constructor to work with Hadoop {@link org.apache.hadoop.conf.Configuration}.
    */
   public TaskExecutor(Configuration conf) {
     this(conf.getInt(ConfigurationKeys.TASK_EXECUTOR_THREADPOOL_SIZE_KEY,
             ConfigurationKeys.DEFAULT_TASK_EXECUTOR_THREADPOOL_SIZE),
         conf.getInt(ConfigurationKeys.TASK_RETRY_THREAD_POOL_CORE_SIZE_KEY,
             ConfigurationKeys.DEFAULT_TASK_RETRY_THREAD_POOL_CORE_SIZE),
         conf.getLong(ConfigurationKeys.TASK_RETRY_INTERVAL_IN_SEC_KEY,
             ConfigurationKeys.DEFAULT_TASK_RETRY_INTERVAL_IN_SEC),
         conf.getInt(ConfigurationKeys.QUEUED_TASK_TIME_MAX_SIZE,
             ConfigurationKeys.DEFAULT_QUEUED_TASK_TIME_MAX_SIZE),
         conf.getLong(ConfigurationKeys.QUEUED_TASK_TIME_MAX_AGE,
             ConfigurationKeys.DEFAULT_QUEUED_TASK_TIME_MAX_AGE),
         conf.getInt(ConfigurationKeys.METRIC_TIMER_WINDOW_SIZE_IN_MINUTES,
             ConfigurationKeys.DEFAULT_METRIC_TIMER_WINDOW_SIZE_IN_MINUTES));
     Log4jConfigurationHelper.setLogLevel(conf.getTrimmedStringCollection(Log4jConfigurationHelper.LOG_LEVEL_OVERRIDE_MAP));
   }

   @Override
   protected void startUp()
       throws Exception {
     LOG.info("Starting the task executor");
     if (this.taskExecutor.isShutdown() || this.taskExecutor.isTerminated()) {
       throw new IllegalStateException("Task thread pool executor is shutdown or terminated");
     }
     if (this.forkExecutor.isShutdown() || this.forkExecutor.isTerminated()) {
       throw new IllegalStateException("Fork thread pool executor is shutdown or terminated");
     }
   }

   @Override
   protected void shutDown()
       throws Exception {
     LOG.info("Stopping the task executor");
     try {
       ExecutorsUtils.shutdownExecutorService(this.taskExecutor, Optional.of(LOG));
     } finally {
       ExecutorsUtils.shutdownExecutorService(this.forkExecutor, Optional.of(LOG));
     }
   }

   /**
    * Execute a {@link Task}.
    *
    * @param task {@link Task} to be executed
    */
   public void execute(Task task) {
     LOG.info(String.format("Executing task %s", task.getTaskId()));
     this.taskExecutor.execute(new TrackingTask(task));
   }

   /**
    * Submit a {@link Task} to run.
    *
    * @param task {@link Task} to be submitted
    * @return a {@link java.util.concurrent.Future} for the submitted {@link Task}
    */
   public Future<?> submit(Task task) {
     LOG.info(String.format("Submitting task %s", task.getTaskId()));
     return this.taskExecutor.submit(new TrackingTask(task));
   }

   /**
    * Execute a {@link Fork}.
    *
    * @param fork {@link Fork} to be executed
    */
   public void execute(Fork fork) {
     LOG.info(String.format("Executing fork %d of task %s", fork.getIndex(), fork.getTaskId()));
     this.forkExecutor.execute(fork);
   }

   /**
    * Submit a {@link Fork} to run.
    *
    * @param fork {@link Fork} to be submitted
    * @return a {@link java.util.concurrent.Future} for the submitted {@link Fork}
    */
   public Future<?> submit(Fork fork) {
     LOG.info(String.format("Submitting fork %d of task %s", fork.getIndex(), fork.getTaskId()));
     return this.forkExecutor.submit(fork);
   }

   /**
    * Retry a failed {@link Task}.
    *
    * @param task failed {@link Task} to be retried
    */
   public void retry(Task task) {
     if (GobblinMetrics.isEnabled(task.getTaskState().getWorkunit()) &&
         task.getTaskState().contains(ConfigurationKeys.FORK_BRANCHES_KEY)) {
       // Adjust metrics to clean up numbers from the failed task
       task.getTaskState()
           .adjustJobMetricsOnRetry(task.getTaskState().getPropAsInt(ConfigurationKeys.FORK_BRANCHES_KEY));
     }

     // Task retry interval increases linearly with number of retries
     long interval = task.getRetryCount() * this.retryIntervalInSeconds;
     // Schedule the retry of the failed task
     this.taskExecutor.schedule(new TrackingTask(task, interval, TimeUnit.SECONDS), interval, TimeUnit.SECONDS);
     LOG.info(String.format("Scheduled retry of failed task %s to run in %d seconds", task.getTaskId(), interval));
     task.incrementRetryCount();
   }

   public MetricSet getTaskExecutorQueueMetricSet() {
     return this.metricSet;
   }

   private synchronized void calculateMetrics() {
     long currentTimeMillis = System.currentTimeMillis();
     if (lastCalculationTime < currentTimeMillis - TimeUnit.SECONDS.toMillis(10)) {
       LOG.debug("Starting metric calculation.");
       int currentQueuedTaskCount = 0;
       int futureQueuedTaskCount = 0;
       long currentQueuedTaskTotalTime = 0;
       for (Map.Entry<String, Long> queuedTask : this.queuedTasks.entrySet()) {
         if (queuedTask.getValue() <= currentTimeMillis) {
           currentQueuedTaskCount++;
           long currentQueuedTaskTime = currentTimeMillis - queuedTask.getValue();
           currentQueuedTaskTotalTime += currentQueuedTaskTime;
           LOG.debug(String.format("Task %s has been waiting in the queue for %d ms.", queuedTask.getKey(), currentQueuedTaskTime));
         } else {
           futureQueuedTaskCount++;
         }
       }
       if (futureQueuedTaskCount > 0) {
         LOG.debug(String.format("%d tasks were ignored during metric calculations because they are scheduled to run in the future.", futureQueuedTaskCount));
       }

       this.currentQueuedTaskCount.set(currentQueuedTaskCount);
       this.currentQueuedTaskTotalTime.set(currentQueuedTaskTotalTime);
       LOG.debug(String.format("%d current tasks have been waiting for a total of %d ms.", currentQueuedTaskCount, currentQueuedTaskTotalTime));

       int historicalQueuedTaskCount = 0;
       long historicalQueuedTaskTotalTime = 0;
       long cutoff = currentTimeMillis - queuedTaskTimeMaxAge;
       Iterator<Map.Entry<Long, Long>> iterator = queuedTaskTimeHistorical.descendingMap().entrySet().iterator();
       while (iterator.hasNext()) {
         try {
           Map.Entry<Long, Long> historicalQueuedTask = iterator.next();
           if (historicalQueuedTask.getKey() < cutoff || historicalQueuedTaskCount >= queuedTaskTimeMaxSize) {
             LOG.debug(String.format("Task started at %d is before the cutoff of %d and is being removed. Queue time %d will be removed from metric calculations.", historicalQueuedTask.getKey(), cutoff, historicalQueuedTask.getValue()));
             iterator.remove();
           } else {
             historicalQueuedTaskCount++;
             historicalQueuedTaskTotalTime += historicalQueuedTask.getValue();
             LOG.debug(String.format("Task started at %d is after cutoff. Queue time %d will be used in metric calculations.", historicalQueuedTask.getKey(), historicalQueuedTask.getValue()));
           }
         } catch (NoSuchElementException e) {
           LOG.warn("Ran out of items in historical task queue time set.");
         }
       }
       this.historicalQueuedTaskCount.set(historicalQueuedTaskCount);
       this.historicalQueuedTaskTotalTime.set(historicalQueuedTaskTotalTime);
       LOG.debug(String.format("%d historical tasks have been waiting for a total of %d ms.", historicalQueuedTaskCount, historicalQueuedTaskTotalTime));

       int totalQueuedTaskCount = currentQueuedTaskCount + historicalQueuedTaskCount;
       long totalQueuedTaskTime = currentQueuedTaskTotalTime + historicalQueuedTaskTotalTime;
       this.queuedTaskCount.set(totalQueuedTaskCount);
       this.queuedTaskTotalTime.set(totalQueuedTaskTime);
       LOG.debug(String.format("%d tasks have been waiting for a total of %d ms.", totalQueuedTaskCount, totalQueuedTaskTime));

       this.lastCalculationTime = currentTimeMillis;
       LOG.debug("Finished metric calculation.");
     } else {
       LOG.debug("Skipped metric calculation because not enough time has elapsed since the last calculation.");
     }
   }

   private class TaskExecutorQueueMetricSet implements MetricSet {
     @Override
     public Map<String, Metric> getMetrics() {
       final Map<String, Metric> metrics = new HashMap<>();
       metrics.put(name("queued", "current", "count"), new Gauge<Integer>() {
         @Override
         public Integer getValue() {
           calculateMetrics();
           return currentQueuedTaskCount.intValue();
         }
       });
       metrics.put(name("queued", "historical", "count"), new Gauge<Integer>() {
         @Override
         public Integer getValue() {
           calculateMetrics();
           return historicalQueuedTaskCount.intValue();
         }
       });
       metrics.put(name("queued", "count"), new Gauge<Integer>() {
         @Override
         public Integer getValue() {
           calculateMetrics();
           return queuedTaskCount.intValue();
         }
       });
       metrics.put(name("queued", "current", "time", "total"), new Gauge<Long>() {
         @Override
         public Long getValue() {
           calculateMetrics();
           return currentQueuedTaskTotalTime.longValue();
         }
       });
       metrics.put(name("queued", "historical", "time", "total"), new Gauge<Long>() {
         @Override
         public Long getValue() {
           calculateMetrics();
           return historicalQueuedTaskTotalTime.longValue();
         }
       });
       metrics.put(name("queued", "time", "total"), new Gauge<Long>() {
         @Override
         public Long getValue() {
           calculateMetrics();
           return queuedTaskTotalTime.longValue();
         }
       });
       metrics.put(name("running", "count"), runningTaskCount);
       metrics.put(name("successful", "count"), successfulTaskCount);
       metrics.put(name("failed", "count"), failedTaskCount);
       return Collections.unmodifiableMap(metrics);
     }
   }

   private class TrackingTask implements Runnable {
     private Task underlyingTask;

     public TrackingTask(Task task) {
       this(task, 0, TimeUnit.SECONDS);
     }

     public TrackingTask(Task task, long interval, TimeUnit timeUnit) {
       long now = System.currentTimeMillis();
       long timeToRun = now + timeUnit.toMillis(interval);
       LOG.debug(String.format("Task %s queued to run %s.", task.getTaskId(), timeToRun <= now ? "now" : "at " + timeToRun));
       queuedTasks.putIfAbsent(task.getTaskId(), timeToRun);
       this.underlyingTask = task;
     }

     @Override
     public void run() {
       long startTime = System.currentTimeMillis();
       onStart(startTime);
       try {
         this.underlyingTask.run();
         successfulTaskCount.mark();
       } catch (Exception e) {
         failedTaskCount.mark();
         LOG.error(String.format("Task %s failed", underlyingTask.getTaskId()), e);
         throw e;
       } finally {
         runningTaskCount.dec();
       }
     }

     private void onStart(long startTime) {
       Long queueTime = queuedTasks.remove(this.underlyingTask.getTaskId());
       long workUnitCreationTime = this.underlyingTask.getTaskContext().getTaskState().getPropAsLong(ConfigurationKeys.WORK_UNIT_CREATION_TIME_IN_MILLIS, 0);
       long timeInQueue = startTime - queueTime;
       long timeSinceWorkUnitCreation = startTime - workUnitCreationTime;

       taskCreateAndRunTimer.update(timeSinceWorkUnitCreation, TimeUnit.MILLISECONDS);

       LOG.debug(String.format("Task %s started. Saving queued time of %d ms to history.", underlyingTask.getTaskId(), timeInQueue));
       queuedTaskTimeHistorical.putIfAbsent(System.currentTimeMillis(), timeInQueue);
       runningTaskCount.inc();
     }
   }
 }
	/*
	* Licensed to the Apache Software Foundation (ASF) under one or more
	* contributor license agreements. See the NOTICE file distributed with
	* this work for additional information regarding copyright ownership.
	* The ASF licenses this file to You 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.gobblin.runtime;

	import java.util.Collections;
	import java.util.HashMap;
	import java.util.Iterator;
	import java.util.Map;
	import java.util.NoSuchElementException;
	import java.util.Properties;
	import java.util.concurrent.ConcurrentSkipListMap;
	import java.util.concurrent.ExecutorService;
	import java.util.concurrent.Executors;
	import java.util.concurrent.Future;
	import java.util.concurrent.ScheduledExecutorService;
	import java.util.concurrent.SynchronousQueue;
	import java.util.concurrent.ThreadPoolExecutor;
	import java.util.concurrent.TimeUnit;
	import java.util.concurrent.atomic.AtomicInteger;
	import java.util.concurrent.atomic.AtomicLong;

	import com.codahale.metrics.Counter;
	import com.codahale.metrics.Gauge;
	import com.codahale.metrics.Meter;
	import com.codahale.metrics.Metric;
	import com.codahale.metrics.MetricSet;
	import com.codahale.metrics.SlidingTimeWindowReservoir;
	import com.codahale.metrics.Timer;
	import com.google.common.base.Optional;
	import com.google.common.base.Preconditions;
	import com.google.common.collect.Maps;
	import com.google.common.util.concurrent.AbstractIdleService;

	import org.apache.hadoop.conf.Configuration;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	import org.apache.gobblin.configuration.ConfigurationKeys;
	import org.apache.gobblin.metrics.GobblinMetrics;
	import org.apache.gobblin.runtime.fork.Fork;
	import org.apache.gobblin.util.ExecutorsUtils;
	import org.apache.gobblin.util.logs.Log4jConfigurationHelper;

	import lombok.Getter;

	import static com.codahale.metrics.MetricRegistry.name;


	/**
	* A class for executing {@link Task}s and retrying failed ones as well as for executing {@link Fork}s.
	*
	* @author Yinan Li
	*/
	public class TaskExecutor extends AbstractIdleService {

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

	// Thread pool executor for running tasks
	private final ScheduledExecutorService taskExecutor;

	// A separate thread pool executor for running forks of tasks
	@Getter
	private final ExecutorService forkExecutor;

	// Task retry interval
	@Getter
	private final long retryIntervalInSeconds;

	// The maximum number of items in the queued task time map.
	@Getter
	private final int queuedTaskTimeMaxSize;

	// The maximum age of the items in the queued task time map.
	@Getter
	private final long queuedTaskTimeMaxAge ;

	// Map of queued task ids to queue times. The key is the task id, the value is the time the task was queued. If the
	// task is being retried, the time may be in the future. Entries with time in the future will not be counted as
	// queued until the time is in the past.
	private final Map<String, Long> queuedTasks = Maps.newConcurrentMap();

	// Set of historical queued task times. The key is the UTC epoch time the task started, the value is the milliseconds
	// the task waited to start.
	private final ConcurrentSkipListMap<Long, Long> queuedTaskTimeHistorical = new ConcurrentSkipListMap<>();

	// The timestamp for the last time the metrics were calculated.
	private long lastCalculationTime = 0;

	// The total number of tasks currently queued and queued over the historical lookback period.
	@Getter
	private AtomicInteger queuedTaskCount = new AtomicInteger();

	// The total number of tasks currently queued.
	@Getter
	private AtomicInteger currentQueuedTaskCount = new AtomicInteger();

	// The total number of tasks queued over the historical lookback period.
	@Getter
	private AtomicInteger historicalQueuedTaskCount = new AtomicInteger();

	// The total time tasks have currently been in the queue and were in the queue during the historical lookback period.
	@Getter
	private AtomicLong queuedTaskTotalTime = new AtomicLong();

	// The total time tasks have currently been in the queue.
	@Getter
	private AtomicLong currentQueuedTaskTotalTime = new AtomicLong();

	// The total time tasks have been in the queue during the historical lookback period.
	@Getter
	private AtomicLong historicalQueuedTaskTotalTime = new AtomicLong();

	// Count of running tasks.
	@Getter
	private final Counter runningTaskCount = new Counter();

	// Count of failed tasks.
	@Getter
	private final Meter successfulTaskCount = new Meter();

	// Count of failed tasks.
	@Getter
	private final Meter failedTaskCount = new Meter();

	@Getter
	private final Timer taskCreateAndRunTimer;

	// The metric set exposed from the task executor.
	private final TaskExecutorQueueMetricSet metricSet = new TaskExecutorQueueMetricSet();

	/**
	* Constructor used internally.
	*/
	private TaskExecutor(int taskExecutorThreadPoolSize, int coreRetryThreadPoolSize, long retryIntervalInSeconds,
	int queuedTaskTimeMaxSize, long queuedTaskTimeMaxAge, int timerWindowSize) {
	Preconditions.checkArgument(taskExecutorThreadPoolSize > 0, "Task executor thread pool size should be positive");
	Preconditions.checkArgument(retryIntervalInSeconds > 0, "Task retry interval should be positive");
	Preconditions.checkArgument(queuedTaskTimeMaxSize > 0, "Queued task time max size should be positive");
	Preconditions.checkArgument(queuedTaskTimeMaxAge > 0, "Queued task time max age should be positive");

	// Currently a fixed-size thread pool is used to execute tasks. We probably need to revisit this later.
	this.taskExecutor = ExecutorsUtils.loggingDecorator(Executors.newScheduledThreadPool(
	taskExecutorThreadPoolSize,
	ExecutorsUtils.newThreadFactory(Optional.of(LOG), Optional.of("TaskExecutor-%d"))));

	this.retryIntervalInSeconds = retryIntervalInSeconds;
	this.queuedTaskTimeMaxSize = queuedTaskTimeMaxSize;
	this.queuedTaskTimeMaxAge = queuedTaskTimeMaxAge;
	this.taskCreateAndRunTimer = new Timer(new SlidingTimeWindowReservoir(timerWindowSize, TimeUnit.MINUTES));

	this.forkExecutor = ExecutorsUtils.loggingDecorator(
	new ThreadPoolExecutor(
	// The core thread pool size is equal to that of the task
	// executor as there's at least one fork per task
	taskExecutorThreadPoolSize,
	// The fork executor thread pool size is essentially unbounded. This is to make sure all forks of
	// a task get a thread to run so all forks of the task are making progress. This is necessary since
	// otherwise the parent task will be blocked if the record queue (bounded) of some fork is full and
	// that fork has not yet started to run because of no available thread. The task cannot proceed in
	// this case because it has to make sure every records go to every forks.
	Integer.MAX_VALUE,
	0L,
	TimeUnit.MILLISECONDS,
	// The work queue is a SynchronousQueue. This essentially forces a new thread to be created for each fork.
	new SynchronousQueue<Runnable>(),
	ExecutorsUtils.newThreadFactory(Optional.of(LOG), Optional.of("ForkExecutor-%d"))));
	}

	/**
	* Constructor to work with {@link java.util.Properties}.
	*/
	public TaskExecutor(Properties properties) {
	this(Integer.parseInt(properties.getProperty(ConfigurationKeys.TASK_EXECUTOR_THREADPOOL_SIZE_KEY,
	Integer.toString(ConfigurationKeys.DEFAULT_TASK_EXECUTOR_THREADPOOL_SIZE))),
	Integer.parseInt(properties.getProperty(ConfigurationKeys.TASK_RETRY_THREAD_POOL_CORE_SIZE_KEY,
	Integer.toString(ConfigurationKeys.DEFAULT_TASK_RETRY_THREAD_POOL_CORE_SIZE))),
	Long.parseLong(properties.getProperty(ConfigurationKeys.TASK_RETRY_INTERVAL_IN_SEC_KEY,
	Long.toString(ConfigurationKeys.DEFAULT_TASK_RETRY_INTERVAL_IN_SEC))),
	Integer.parseInt(properties.getProperty(ConfigurationKeys.QUEUED_TASK_TIME_MAX_SIZE,
	Integer.toString(ConfigurationKeys.DEFAULT_QUEUED_TASK_TIME_MAX_SIZE))),
	Long.parseLong(properties.getProperty(ConfigurationKeys.QUEUED_TASK_TIME_MAX_AGE,
	Long.toString(ConfigurationKeys.DEFAULT_QUEUED_TASK_TIME_MAX_AGE))),
	Integer.parseInt(properties.getProperty(ConfigurationKeys.METRIC_TIMER_WINDOW_SIZE_IN_MINUTES,
	Integer.toString(ConfigurationKeys.DEFAULT_METRIC_TIMER_WINDOW_SIZE_IN_MINUTES))));
	}

	/**
	* Constructor to work with Hadoop {@link org.apache.hadoop.conf.Configuration}.
	*/
	public TaskExecutor(Configuration conf) {
	this(conf.getInt(ConfigurationKeys.TASK_EXECUTOR_THREADPOOL_SIZE_KEY,
	ConfigurationKeys.DEFAULT_TASK_EXECUTOR_THREADPOOL_SIZE),
	conf.getInt(ConfigurationKeys.TASK_RETRY_THREAD_POOL_CORE_SIZE_KEY,
	ConfigurationKeys.DEFAULT_TASK_RETRY_THREAD_POOL_CORE_SIZE),
	conf.getLong(ConfigurationKeys.TASK_RETRY_INTERVAL_IN_SEC_KEY,
	ConfigurationKeys.DEFAULT_TASK_RETRY_INTERVAL_IN_SEC),
	conf.getInt(ConfigurationKeys.QUEUED_TASK_TIME_MAX_SIZE,
	ConfigurationKeys.DEFAULT_QUEUED_TASK_TIME_MAX_SIZE),
	conf.getLong(ConfigurationKeys.QUEUED_TASK_TIME_MAX_AGE,
	ConfigurationKeys.DEFAULT_QUEUED_TASK_TIME_MAX_AGE),
	conf.getInt(ConfigurationKeys.METRIC_TIMER_WINDOW_SIZE_IN_MINUTES,
	ConfigurationKeys.DEFAULT_METRIC_TIMER_WINDOW_SIZE_IN_MINUTES));
	Log4jConfigurationHelper.setLogLevel(conf.getTrimmedStringCollection(Log4jConfigurationHelper.LOG_LEVEL_OVERRIDE_MAP));
	}

	@Override
	protected void startUp()
	throws Exception {
	LOG.info("Starting the task executor");
	if (this.taskExecutor.isShutdown() \|\| this.taskExecutor.isTerminated()) {
	throw new IllegalStateException("Task thread pool executor is shutdown or terminated");
	}
	if (this.forkExecutor.isShutdown() \|\| this.forkExecutor.isTerminated()) {
	throw new IllegalStateException("Fork thread pool executor is shutdown or terminated");
	}
	}

	@Override
	protected void shutDown()
	throws Exception {
	LOG.info("Stopping the task executor");
	try {
	ExecutorsUtils.shutdownExecutorService(this.taskExecutor, Optional.of(LOG));
	} finally {
	ExecutorsUtils.shutdownExecutorService(this.forkExecutor, Optional.of(LOG));
	}
	}

	/**
	* Execute a {@link Task}.
	*
	* @param task {@link Task} to be executed
	*/
	public void execute(Task task) {
	LOG.info(String.format("Executing task %s", task.getTaskId()));
	this.taskExecutor.execute(new TrackingTask(task));
	}

	/**
	* Submit a {@link Task} to run.
	*
	* @param task {@link Task} to be submitted
	* @return a {@link java.util.concurrent.Future} for the submitted {@link Task}
	*/
	public Future<?> submit(Task task) {
	LOG.info(String.format("Submitting task %s", task.getTaskId()));
	return this.taskExecutor.submit(new TrackingTask(task));
	}

	/**
	* Execute a {@link Fork}.
	*
	* @param fork {@link Fork} to be executed
	*/
	public void execute(Fork fork) {
	LOG.info(String.format("Executing fork %d of task %s", fork.getIndex(), fork.getTaskId()));
	this.forkExecutor.execute(fork);
	}

	/**
	* Submit a {@link Fork} to run.
	*
	* @param fork {@link Fork} to be submitted
	* @return a {@link java.util.concurrent.Future} for the submitted {@link Fork}
	*/
	public Future<?> submit(Fork fork) {
	LOG.info(String.format("Submitting fork %d of task %s", fork.getIndex(), fork.getTaskId()));
	return this.forkExecutor.submit(fork);
	}

	/**
	* Retry a failed {@link Task}.
	*
	* @param task failed {@link Task} to be retried
	*/
	public void retry(Task task) {
	if (GobblinMetrics.isEnabled(task.getTaskState().getWorkunit()) &&
	task.getTaskState().contains(ConfigurationKeys.FORK_BRANCHES_KEY)) {
	// Adjust metrics to clean up numbers from the failed task
	task.getTaskState()
	.adjustJobMetricsOnRetry(task.getTaskState().getPropAsInt(ConfigurationKeys.FORK_BRANCHES_KEY));
	}

	// Task retry interval increases linearly with number of retries
	long interval = task.getRetryCount() * this.retryIntervalInSeconds;
	// Schedule the retry of the failed task
	this.taskExecutor.schedule(new TrackingTask(task, interval, TimeUnit.SECONDS), interval, TimeUnit.SECONDS);
	LOG.info(String.format("Scheduled retry of failed task %s to run in %d seconds", task.getTaskId(), interval));
	task.incrementRetryCount();
	}

	public MetricSet getTaskExecutorQueueMetricSet() {
	return this.metricSet;
	}

	private synchronized void calculateMetrics() {
	long currentTimeMillis = System.currentTimeMillis();
	if (lastCalculationTime < currentTimeMillis - TimeUnit.SECONDS.toMillis(10)) {
	LOG.debug("Starting metric calculation.");
	int currentQueuedTaskCount = 0;
	int futureQueuedTaskCount = 0;
	long currentQueuedTaskTotalTime = 0;
	for (Map.Entry<String, Long> queuedTask : this.queuedTasks.entrySet()) {
	if (queuedTask.getValue() <= currentTimeMillis) {
	currentQueuedTaskCount++;
	long currentQueuedTaskTime = currentTimeMillis - queuedTask.getValue();
	currentQueuedTaskTotalTime += currentQueuedTaskTime;
	LOG.debug(String.format("Task %s has been waiting in the queue for %d ms.", queuedTask.getKey(), currentQueuedTaskTime));
	} else {
	futureQueuedTaskCount++;
	}
	}
	if (futureQueuedTaskCount > 0) {
	LOG.debug(String.format("%d tasks were ignored during metric calculations because they are scheduled to run in the future.", futureQueuedTaskCount));
	}

	this.currentQueuedTaskCount.set(currentQueuedTaskCount);
	this.currentQueuedTaskTotalTime.set(currentQueuedTaskTotalTime);
	LOG.debug(String.format("%d current tasks have been waiting for a total of %d ms.", currentQueuedTaskCount, currentQueuedTaskTotalTime));

	int historicalQueuedTaskCount = 0;
	long historicalQueuedTaskTotalTime = 0;
	long cutoff = currentTimeMillis - queuedTaskTimeMaxAge;
	Iterator<Map.Entry<Long, Long>> iterator = queuedTaskTimeHistorical.descendingMap().entrySet().iterator();
	while (iterator.hasNext()) {
	try {
	Map.Entry<Long, Long> historicalQueuedTask = iterator.next();
	if (historicalQueuedTask.getKey() < cutoff \|\| historicalQueuedTaskCount >= queuedTaskTimeMaxSize) {
	LOG.debug(String.format("Task started at %d is before the cutoff of %d and is being removed. Queue time %d will be removed from metric calculations.", historicalQueuedTask.getKey(), cutoff, historicalQueuedTask.getValue()));
	iterator.remove();
	} else {
	historicalQueuedTaskCount++;
	historicalQueuedTaskTotalTime += historicalQueuedTask.getValue();
	LOG.debug(String.format("Task started at %d is after cutoff. Queue time %d will be used in metric calculations.", historicalQueuedTask.getKey(), historicalQueuedTask.getValue()));
	}
	} catch (NoSuchElementException e) {
	LOG.warn("Ran out of items in historical task queue time set.");
	}
	}
	this.historicalQueuedTaskCount.set(historicalQueuedTaskCount);
	this.historicalQueuedTaskTotalTime.set(historicalQueuedTaskTotalTime);
	LOG.debug(String.format("%d historical tasks have been waiting for a total of %d ms.", historicalQueuedTaskCount, historicalQueuedTaskTotalTime));

	int totalQueuedTaskCount = currentQueuedTaskCount + historicalQueuedTaskCount;
	long totalQueuedTaskTime = currentQueuedTaskTotalTime + historicalQueuedTaskTotalTime;
	this.queuedTaskCount.set(totalQueuedTaskCount);
	this.queuedTaskTotalTime.set(totalQueuedTaskTime);
	LOG.debug(String.format("%d tasks have been waiting for a total of %d ms.", totalQueuedTaskCount, totalQueuedTaskTime));

	this.lastCalculationTime = currentTimeMillis;
	LOG.debug("Finished metric calculation.");
	} else {
	LOG.debug("Skipped metric calculation because not enough time has elapsed since the last calculation.");
	}
	}

	private class TaskExecutorQueueMetricSet implements MetricSet {
	@Override
	public Map<String, Metric> getMetrics() {
	final Map<String, Metric> metrics = new HashMap<>();
	metrics.put(name("queued", "current", "count"), new Gauge<Integer>() {
	@Override
	public Integer getValue() {
	calculateMetrics();
	return currentQueuedTaskCount.intValue();
	}
	});
	metrics.put(name("queued", "historical", "count"), new Gauge<Integer>() {
	@Override
	public Integer getValue() {
	calculateMetrics();
	return historicalQueuedTaskCount.intValue();
	}
	});
	metrics.put(name("queued", "count"), new Gauge<Integer>() {
	@Override
	public Integer getValue() {
	calculateMetrics();
	return queuedTaskCount.intValue();
	}
	});
	metrics.put(name("queued", "current", "time", "total"), new Gauge<Long>() {
	@Override
	public Long getValue() {
	calculateMetrics();
	return currentQueuedTaskTotalTime.longValue();
	}
	});
	metrics.put(name("queued", "historical", "time", "total"), new Gauge<Long>() {
	@Override
	public Long getValue() {
	calculateMetrics();
	return historicalQueuedTaskTotalTime.longValue();
	}
	});
	metrics.put(name("queued", "time", "total"), new Gauge<Long>() {
	@Override
	public Long getValue() {
	calculateMetrics();
	return queuedTaskTotalTime.longValue();
	}
	});
	metrics.put(name("running", "count"), runningTaskCount);
	metrics.put(name("successful", "count"), successfulTaskCount);
	metrics.put(name("failed", "count"), failedTaskCount);
	return Collections.unmodifiableMap(metrics);
	}
	}

	private class TrackingTask implements Runnable {
	private Task underlyingTask;

	public TrackingTask(Task task) {
	this(task, 0, TimeUnit.SECONDS);
	}

	public TrackingTask(Task task, long interval, TimeUnit timeUnit) {
	long now = System.currentTimeMillis();
	long timeToRun = now + timeUnit.toMillis(interval);
	LOG.debug(String.format("Task %s queued to run %s.", task.getTaskId(), timeToRun <= now ? "now" : "at " + timeToRun));
	queuedTasks.putIfAbsent(task.getTaskId(), timeToRun);
	this.underlyingTask = task;
	}

	@Override
	public void run() {
	long startTime = System.currentTimeMillis();
	onStart(startTime);
	try {
	this.underlyingTask.run();
	successfulTaskCount.mark();
	} catch (Exception e) {
	failedTaskCount.mark();
	LOG.error(String.format("Task %s failed", underlyingTask.getTaskId()), e);
	throw e;
	} finally {
	runningTaskCount.dec();
	}
	}

	private void onStart(long startTime) {
	Long queueTime = queuedTasks.remove(this.underlyingTask.getTaskId());
	long workUnitCreationTime = this.underlyingTask.getTaskContext().getTaskState().getPropAsLong(ConfigurationKeys.WORK_UNIT_CREATION_TIME_IN_MILLIS, 0);
	long timeInQueue = startTime - queueTime;
	long timeSinceWorkUnitCreation = startTime - workUnitCreationTime;

	taskCreateAndRunTimer.update(timeSinceWorkUnitCreation, TimeUnit.MILLISECONDS);

	LOG.debug(String.format("Task %s started. Saving queued time of %d ms to history.", underlyingTask.getTaskId(), timeInQueue));
	queuedTaskTimeHistorical.putIfAbsent(System.currentTimeMillis(), timeInQueue);
	runningTaskCount.inc();
	}
	}
	}