001package runwar.util;
002import java.util.concurrent.ScheduledThreadPoolExecutor;
003import java.util.concurrent.locks.Condition;
004import java.util.concurrent.locks.ReentrantLock;
005
006public class PausableScheduledThreadPoolExecutor extends ScheduledThreadPoolExecutor {      
007  public boolean isPaused;
008  private ReentrantLock pauseLock = new ReentrantLock();
009  private Condition unpaused = pauseLock.newCondition();
010  private Latch activeTasksLatch = new Latch();
011
012  private class Latch {
013    private final Object synchObj = new Object();
014    private int count;
015
016    public boolean awaitZero(long waitMS) throws InterruptedException {
017      long startTime = System.currentTimeMillis();
018      synchronized (synchObj) {
019        while (count > 0) {
020          if ( waitMS != 0) {
021            synchObj.wait(waitMS);
022            long curTime = System.currentTimeMillis();
023            if ( (curTime - startTime) > waitMS ) {                
024              return count <= 0;
025            }
026          }
027          else
028            synchObj.wait();
029        }
030        return count <= 0; 
031      }
032    }
033    public void countDown() {
034      synchronized (synchObj) {
035        if (--count <= 0) {
036          // assert count >= 0;              
037          synchObj.notifyAll();
038        }
039      }
040    }
041    public void countUp() {
042      synchronized (synchObj) {
043        count++;
044      }
045    }    
046  }
047
048  /**
049   * Default constructor for a simple fixed threadpool
050   */
051  public PausableScheduledThreadPoolExecutor(int corePoolSize) {
052    super(corePoolSize);
053  }
054
055  /**
056   * Executed before a task is assigned to a thread.
057   */
058  @Override
059  protected void beforeExecute(Thread t, Runnable r) {
060    pauseLock.lock();
061    try {
062      while (isPaused)
063        unpaused.await();
064    } catch (InterruptedException ie) {
065      t.interrupt();
066    } finally {
067      pauseLock.unlock();
068    }
069
070    activeTasksLatch.countUp();
071    super.beforeExecute(t, r);
072  }
073
074  @Override
075  protected void afterExecute(Runnable r, Throwable t) {
076    try {
077      super.afterExecute(r, t);
078    }
079    finally {
080      activeTasksLatch.countDown();
081    }
082  }
083
084  /**
085   * Pause the threadpool. Running tasks will continue running, but new tasks
086   * will not start untill the threadpool is resumed.
087   */
088  public void pause() {
089    pauseLock.lock();
090    try {
091      isPaused = true;
092    } finally {
093      pauseLock.unlock();
094    }
095  }
096
097  /**
098   * Wait for all active tasks to end.
099   */ 
100  public boolean await(long timeoutMS) {
101    // assert isPaused;
102    try {
103      return activeTasksLatch.awaitZero(timeoutMS);
104    } catch (InterruptedException e) {
105      // log e, or rethrow maybe
106    }
107    return false;
108  }
109
110  /**
111   * Resume the threadpool.
112   */
113  public void resume() {
114    pauseLock.lock();
115    try {
116      isPaused = false;
117      unpaused.signalAll();
118    } finally {
119      pauseLock.unlock();
120    }
121  }
122
123}