/*
  * JBoss, Home of Professional Open Source
  * Copyright 2005, JBoss Inc., and individual contributors as indicated
  * by the @authors tag. See the copyright.txt in the distribution for a
  * full listing of individual contributors.
  *
  * This 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 (at your option) any later version.
  *
  * This software 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 software; if not, write to the Free
  * Software Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA
  * 02110-1301 USA, or see the FSF site: http://www.fsf.org.
  */
package org.jboss.util.threadpool;

import java.security.AccessController;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Map;
import java.util.PriorityQueue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;

import org.jboss.logging.Logger;
import org.jboss.util.collection.WeakValueHashMap;
import org.jboss.util.loading.ClassLoaderSource;
import org.jboss.util.loading.ContextClassLoaderSwitcher;


/**
 * A basic thread pool.
 * TODO: this port to jdk concurrent still needs to be tested.
 *
 * @author <a href="mailto:adrian@jboss.org">Adrian Brock</a>
 * @author Scott.Stark@jboss.org
 * @version $Revision: 2892 $
 */
@SuppressWarnings("unchecked")
public class BasicThreadPool implements ThreadPool, BasicThreadPoolMBean
{
   // Constants -----------------------------------------------------

   /** The jboss thread group */
   private static final ThreadGroup JBOSS_THREAD_GROUP = new ThreadGroup("JBoss Pooled Threads");

   /** The thread groups */
   private static final Map threadGroups = Collections.synchronizedMap(new WeakValueHashMap());

   /** The internal pool number */
   private static final AtomicInteger lastPoolNumber = new AtomicInteger(0);

   private static Logger log = Logger.getLogger(BasicThreadPool.class);

   // Attributes ----------------------------------------------------

   /** The thread pool name */
   private String name;

   /** The internal pool number */
   private int poolNumber;

   /** The blocking mode */
   private BlockingMode blockingMode = BlockingMode.ABORT;

   /** The pooled executor */
   private ThreadPoolExecutor executor;

   /** The queue */
   private LinkedBlockingQueue queue;

   /** The thread group */
   private ThreadGroup threadGroup;
   
   /** Source for the thread contrext classloader to assign to threads */
   private ClassLoaderSource classLoaderSource;

   private ContextClassLoaderSwitcher classLoaderSwitcher;
   
   /** The last thread number */
   private AtomicInteger lastThreadNumber = new AtomicInteger(0);

   /** Has the pool been stopped? */
   private AtomicBoolean stopped = new AtomicBoolean(false);
   /** The Heap<TimeoutInfo> of tasks ordered by their completion timeout */
   private PriorityQueue<TimeoutInfo> tasksWithTimeouts = new PriorityQueue<TimeoutInfo>(13);
   /** The task completion timeout monitor runnable */
   private TimeoutMonitor timeoutTask;
   /** The trace level logging flag */
   private boolean trace;

   // Static --------------------------------------------------------

   // Constructors --------------------------------------------------

   /**
    * Create a new thread pool
    */
   public BasicThreadPool()
   {
      this("ThreadPool");
   }

   /**
    * Create a new thread pool with a default queue size of 1024, max pool
    * size of 100, min pool size of 4, and a keep alive of 60 seconds.
    *
    * @param name the pool name
    */
   public BasicThreadPool(String name)
   {
      this(name, JBOSS_THREAD_GROUP);
   }

   /**
    * Create a new thread pool with a default queue size of 1024, max pool
    * size of 100, min pool size of 4, and a keep alive of 60 seconds.
    *
    * @param name the pool name
    * @param threadGroup threadGroup
    */
   public BasicThreadPool(String name, ThreadGroup threadGroup)
   {
      trace = log.isTraceEnabled();
      ThreadFactory factory = new ThreadPoolThreadFactory();

      queue = new LinkedBlockingQueue(1024);

      
      executor = new RestoreTCCLThreadPoolExecutor(4, 4, 60, TimeUnit.SECONDS, queue);
      executor.setThreadFactory(factory);
      executor.setRejectedExecutionHandler(new ThreadPoolExecutor.AbortPolicy());
      
      poolNumber = lastPoolNumber.incrementAndGet();
      setName(name);
      this.threadGroup = threadGroup;
   }

   // Public --------------------------------------------------------

   // ThreadPool ----------------------------------------------------

   public void stop(boolean immediate)
   {
      log.debug("stop, immediate="+immediate);
      stopped.set(true);
      if (immediate)
         executor.shutdownNow();
      else
         executor.shutdown();
   }

   public void waitForTasks() throws InterruptedException
   {
      executor.awaitTermination(Long.MAX_VALUE, TimeUnit.MILLISECONDS);
   }
   public void waitForTasks(long maxWaitTime) throws InterruptedException
   {
      executor.awaitTermination(maxWaitTime, TimeUnit.MILLISECONDS);
   }

   public void runTaskWrapper(TaskWrapper wrapper)
   {
      if( trace )
         log.trace("runTaskWrapper, wrapper="+wrapper);
      if (stopped.get())
      {
         wrapper.rejectTask(new ThreadPoolStoppedException("Thread pool has been stopped"));
         return;
      }

      wrapper.acceptTask();

      long completionTimeout = wrapper.getTaskCompletionTimeout();
      TimeoutInfo info;
      if(completionTimeout > 0)
      {
         checkTimeoutMonitor();
         // Install the task in the
         info = new TimeoutInfo(wrapper, completionTimeout);
         tasksWithTimeouts.add(info);
      }
      int waitType = wrapper.getTaskWaitType();
      switch (waitType)
      {
         case Task.WAIT_FOR_COMPLETE:
         {
            executeOnThread(wrapper);
            break;
         }
         default:
         {
            execute(wrapper);
         }
      }
      waitForTask(wrapper);
   }

   public void runTask(Task task)
   {
      BasicTaskWrapper wrapper = new BasicTaskWrapper(task);
      runTaskWrapper(wrapper);
   }

   public void run(Runnable runnable)
   {
      run(runnable, 0, 0);
   }

   public void run(Runnable runnable, long startTimeout, long completeTimeout)
   {
      RunnableTaskWrapper wrapper = new RunnableTaskWrapper(runnable, startTimeout, completeTimeout);
      runTaskWrapper(wrapper);      
   }

   public ThreadGroup getThreadGroup()
   {
      return threadGroup;
   }

   // ThreadPoolMBean implementation --------------------------------

   public String getName()
   {
      return name;
   }

   public void setName(String name)
   {
      this.name = name;
   }

   public int getPoolNumber()
   {
      return poolNumber;
   }

   public String getThreadGroupName()
   {
      return threadGroup.getName();
   }

   public void setThreadGroupName(String threadGroupName)
   {
      ThreadGroup group;
      synchronized(threadGroups)
      {
         group = (ThreadGroup) threadGroups.get(threadGroupName);
         if (group == null)
         {
            group = new ThreadGroup(JBOSS_THREAD_GROUP, threadGroupName);
            threadGroups.put(threadGroupName, group);
         }
      }
      threadGroup = group;
   }

   public int getQueueSize()
   {
      return queue.size();
   }

   public int getMaximumQueueSize()
   {
      return queue.remainingCapacity() + queue.size();
   }

   /**
    * This resets the work queue capacity. This requires recreating the
    * work queue and ThreadPoolExecutor, so this needs to be called
    * before doing any work with the pool.
    *
    * @param size new work queue capacity
    */
   public void setMaximumQueueSize(int size)
   {
      // Reset the executor work queue
      ArrayList tmp = new ArrayList();
      queue.drainTo(tmp);
      queue = new LinkedBlockingQueue(size);
      queue.addAll(tmp);

      ThreadFactory tf = executor.getThreadFactory();
      RejectedExecutionHandler handler = executor.getRejectedExecutionHandler();
      long keepAlive = executor.getKeepAliveTime(TimeUnit.SECONDS);
      int cs = executor.getCorePoolSize();
      int mcs = executor.getMaximumPoolSize();
      executor = new ThreadPoolExecutor(cs, mcs, keepAlive, TimeUnit.SECONDS, queue);
      executor.setThreadFactory(tf);
      executor.setRejectedExecutionHandler(handler);
   }

   public int getPoolSize()
   {
      return executor.getPoolSize();
   }

   public int getMinimumPoolSize()
   {
      return executor.getCorePoolSize();
   }

   public void setMinimumPoolSize(int size)
   {
      synchronized (executor)
      {
         // Don't let the min size > max size
         if (executor.getMaximumPoolSize() < size)
         {
            executor.setCorePoolSize(size);
            executor.setMaximumPoolSize(size);
         }
      }
   }

   public int getMaximumPoolSize()
   {
      return executor.getMaximumPoolSize();
   }
   
   public void setMaximumPoolSize(int size)
   {
      synchronized (executor)
      {
         executor.setCorePoolSize(size);
         executor.setMaximumPoolSize(size);
      }
   }

   public long getKeepAliveTime()
   {
      return executor.getKeepAliveTime(TimeUnit.MILLISECONDS);
   }

   public void setKeepAliveTime(long time)
   {
      executor.setKeepAliveTime(time, TimeUnit.MILLISECONDS);
   }

   public BlockingMode getBlockingMode()
   {
      return blockingMode;
   }

   public void setBlockingMode(BlockingMode mode)
   {
      blockingMode = mode;
      
      if( blockingMode == BlockingMode.RUN )
      {
         executor.setRejectedExecutionHandler(new ThreadPoolExecutor.CallerRunsPolicy());
      }
      else if( blockingMode == BlockingMode.WAIT )
      {
         executor.setRejectedExecutionHandler(new ThreadPoolExecutor.CallerRunsPolicy());
      }
      else if( blockingMode == BlockingMode.DISCARD )
      {
         executor.setRejectedExecutionHandler(new ThreadPoolExecutor.DiscardPolicy());
      }
      else if( blockingMode == BlockingMode.DISCARD_OLDEST )
      {
         executor.setRejectedExecutionHandler(new ThreadPoolExecutor.DiscardOldestPolicy());
      }
      else if( blockingMode == BlockingMode.ABORT )
      {
         executor.setRejectedExecutionHandler(new ThreadPoolExecutor.AbortPolicy());
      }
      else
      {
         throw new IllegalArgumentException("Failed to recognize mode: "+mode);
      }
   }

   /**
    * For backward compatibility with the previous string based mode
    * @param name - the string form of the mode enum
    */ 
   public void setBlockingMode(String name)
   {
      blockingMode = BlockingMode.toBlockingMode(name);
      if( blockingMode == null )
         blockingMode = BlockingMode.ABORT;
   }

   /**
    * For backward compatibility with the previous string based mode
    * This is needed for microcontainer as it gets confused with overloaded
    * setters.
    * @param name - the string form of the mode enum
    */
   public void setBlockingModeString(String name)
   {
      blockingMode = BlockingMode.toBlockingMode(name);
      if( blockingMode == null )
         blockingMode = BlockingMode.ABORT;
   }

   public ClassLoaderSource getClassLoaderSource()
   {
      return classLoaderSource;
   }

   public void setClassLoaderSource(ClassLoaderSource classLoaderSource)
   {
      if (classLoaderSource == null)
      {
         this.classLoaderSource = null;
         this.classLoaderSwitcher = null;
      }
      else if (classLoaderSwitcher == null)
      {
         try
         {
            this.classLoaderSwitcher = (ContextClassLoaderSwitcher) AccessController.doPrivileged(ContextClassLoaderSwitcher.INSTANTIATOR);
            this.classLoaderSource = classLoaderSource;
         }
         catch (SecurityException e)
         {
            log.error("Cannot manage context classloader for pool threads; " +
                      "Do not have setContextClassLoader permission");
         }
      }
      else
      {
         this.classLoaderSource = classLoaderSource;
      }
   }

   public ThreadPool getInstance()
   {
      return this;
   }

   public void stop()
   {
      stop(false);
   }

   // Object overrides ----------------------------------------------

   public String toString()
   {
      return name + '(' + poolNumber + ')';
   }

   // Package protected ---------------------------------------------

   // Protected -----------------------------------------------------

   /**
    * Execute a task on the same thread
    *
    * @param wrapper the task wrapper
    */
   protected void executeOnThread(TaskWrapper wrapper)
   {
      if( trace )
         log.trace("executeOnThread, wrapper="+wrapper);
      wrapper.run();
   }

   /**
    * Execute a task
    *
    * @param wrapper the task wrapper
    */
   protected void execute(TaskWrapper wrapper)
   {
      if( trace )
         log.trace("execute, wrapper="+wrapper);
      try
      {
         executor.execute(wrapper);
      }
      catch (Throwable t)
      {
         wrapper.rejectTask(new ThreadPoolFullException("Error scheduling work: " + wrapper, t));
      }
   }

   /**
    * Wait for a task
    *
    * @param wrapper the task wrapper
    */
   protected void waitForTask(TaskWrapper wrapper)
   {
      wrapper.waitForTask();
   }

   /**
    * Used to lazily create the task completion timeout thread and monitor
    */ 
   protected synchronized void checkTimeoutMonitor()
   {
      if( timeoutTask == null )
         timeoutTask = new TimeoutMonitor(name, log);      
   }
   protected TimeoutInfo getNextTimeout()
   {
      TimeoutInfo info = null;
      if(this.tasksWithTimeouts.isEmpty() == false)
      {
         info = this.tasksWithTimeouts.remove();
      }
      return info;
   }
   
   protected void setDefaultThreadContextClassLoader(Thread thread)
   {
      if (classLoaderSwitcher != null)
      {
         ClassLoader cl = classLoaderSource == null ? null : classLoaderSource.getClassLoader();
         classLoaderSwitcher.setContextClassLoader(thread, cl);
      }
   }

   // Private -------------------------------------------------------

   // Inner classes -------------------------------------------------

   /**
    * A factory for threads
    */
   private class ThreadPoolThreadFactory implements ThreadFactory
   {
      public Thread newThread(Runnable runnable)
      {
         String threadName = BasicThreadPool.this.toString() + "-" + lastThreadNumber.incrementAndGet();
         Thread thread = new Thread(threadGroup, runnable, threadName);
         thread.setDaemon(true);
         BasicThreadPool.this.setDefaultThreadContextClassLoader(thread);
         return thread;
      }
   }
   
   private class RestoreTCCLThreadPoolExecutor extends ThreadPoolExecutor
   {      
      public RestoreTCCLThreadPoolExecutor(int corePoolSize, int maximumPoolSize, 
                                           long keepAliveTime, TimeUnit unit,
                                           BlockingQueue<Runnable> workQueue)
      {
         super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue);
      }

      @Override
      protected void afterExecute(Runnable r, Throwable t)
      {
         try
         {
            super.afterExecute(r, t);
         }
         finally
         {

            BasicThreadPool.this.setDefaultThreadContextClassLoader(Thread.currentThread());
         }
      }
      
      
      
   }

   /** An encapsulation of a task and its completion timeout
    */ 
   private static class TimeoutInfo implements Comparable
   {
      long start;
      long timeoutMS;
      TaskWrapper wrapper;
      boolean firstStop;
      TimeoutInfo(TaskWrapper wrapper, long timeout)
      {
         this.start = System.currentTimeMillis();
         this.timeoutMS = start + timeout;
         this.wrapper = wrapper;
      }
      public void setTimeout(long timeout)
      {
         this.start = System.currentTimeMillis();
         this.timeoutMS = start + timeout;         
      }
      /** Order TimeoutInfo based on the timestamp at which the task needs to
       * be completed by.
       * @param o a TimeoutInfo
       * @return the diff between this timeoutMS and the argument timeoutMS
       */ 
      public int compareTo(Object o)
      {
         TimeoutInfo ti = (TimeoutInfo) o;
         long to0 = timeoutMS;
         long to1 = ti.timeoutMS;
         return (int) (to0 - to1);
      }
      TaskWrapper getTaskWrapper()
      {
         return wrapper;
      }
      public long getTaskCompletionTimeout()
      {
         return wrapper.getTaskCompletionTimeout();
      }
      /** Get the time remaining to the complete timeout timestamp in MS.
       * @param now - the current System.currentTimeMillis value
       * @return the time remaining to the complete timeout timestamp in MS.
       */ 
      public long getTaskCompletionTimeout(long now)
      {
         return timeoutMS - now;
      }
      /** Invoke stopTask on the wrapper and indicate whether this was the first
       * time the task has been notified to stop.
       * @return true if this is the first stopTask, false on the second.
       */ 
      public boolean stopTask()
      {
         wrapper.stopTask();
         boolean wasFirstStop = firstStop == false;
         firstStop = true;
         return wasFirstStop;
      }
   }
   /**
    * The monitor runnable which validates that threads are completing within
    * the task completion timeout limits.
    */ 
   private class TimeoutMonitor implements Runnable
   {
      @SuppressWarnings("hiding")
      final Logger log;
      TimeoutMonitor(String name, Logger log)
      {
         this.log = log;
         Thread t = new Thread(this, name+" TimeoutMonitor");
         t.setDaemon(true);
         t.start();
      }
      /** The monitor thread loops until the pool is shutdown. It waits for
       * tasks with completion timeouts and sleeps until the next completion
       * timeout and then interrupts the associated task thread, and invokes
       * stopTask on the TaskWrapper. A new timeout check is then inserted with
       * a 1 second timeout to validate that the TaskWrapper has exited the
       * run method. If it has not, then the associated task thread is stopped
       * using the deprecated Thread.stop method since this is the only way to
       * abort a thread that is in spin loop for example.
       * 
       * @todo this is not responsive to new tasks with timeouts smaller than
       * the current shortest completion expiration. We probably should interrupt
       * the thread on each insertion into the timeout heap to ensure better
       * responsiveness.
       */ 
      public void run()
      {
         boolean isStopped = stopped.get();
         while( isStopped == false )
         {
            boolean trace = log.isTraceEnabled();
            try
            {
               TimeoutInfo info = getNextTimeout();
               if( info != null )
               {
                  long now = System.currentTimeMillis();
                  long timeToTimeout = info.getTaskCompletionTimeout(now);
                  if( timeToTimeout > 0 )
                  {
                     if( trace )
                     {
                        log.trace("Will check wrapper="+info.getTaskWrapper()
                           +" after "+timeToTimeout);
                     }
                     Thread.sleep(timeToTimeout);
                  }
                  // Check the status of the task
                  TaskWrapper wrapper = info.getTaskWrapper();
                  if( wrapper.isComplete() == false )
                  {
                     if( trace )
                        log.trace("Failed completion check for wrapper="+wrapper);
                     if( info.stopTask() == true )
                     {
                        // Requeue the TimeoutInfo to see that the task exits run
                        info.setTimeout(1000);
                        tasksWithTimeouts.add(info);
                        if( trace )
                           log.trace("Rescheduled completion check for wrapper="+wrapper);
                     }
                  }
               }
               else
               {
                  Thread.sleep(1000);
               }
            }
            catch(InterruptedException e)
            {
               log.debug("Timeout monitor has been interrupted", e);
            }
            catch(Throwable e)
            {
               log.debug("Timeout monitor saw unexpected error", e);               
            }
            isStopped = stopped.get();            
         }
      }
   }
}
