lib/cpp/src/concurrency/ThreadManager.cc

#include "ThreadManager.h"
#include "Exception.h"
#include "Monitor.h"

#include <boost/shared_ptr.hpp>

#include <assert.h>
#include <queue>
#include <set>

#if defined(DEBUG)
#include <iostream>
#endif //defined(DEBUG)

namespace facebook { namespace thrift { namespace concurrency { 

using namespace boost;


/**
 * ThreadManager class
 * 
 * This class manages a pool of threads. It uses a ThreadFactory to create
 * threads.  It never actually creates or destroys worker threads, rather
 * it maintains statistics on number of idle threads, number of active threads,
 * task backlog, and average wait and service times.
 *
 * @author marc
 * @version $Id:$
 */
class ThreadManager::Impl : public ThreadManager  {

 public:
  Impl() : 
    _workerCount(0),
    _workerMaxCount(0),
    _idleCount(0),
    _state(ThreadManager::UNINITIALIZED) {}

  ~Impl() { stop(); }

  void start();

  void stop();

  const ThreadManager::STATE state() const { return _state; }

  shared_ptr<ThreadFactory> threadFactory() const {
    Synchronized s(_monitor); 
    return _threadFactory;
  }
      
  void threadFactory(shared_ptr<ThreadFactory> value) {  
    Synchronized s(_monitor);
    _threadFactory = value;
  }

  void addWorker(size_t value);
  
  void removeWorker(size_t value);
  
  size_t idleWorkerCount() const { return _idleCount; }

  size_t workerCount() const {
    Synchronized s(_monitor); 
    return _workerCount;
  }
  
  size_t pendingTaskCount() const {
    Synchronized s(_monitor); 
    return _tasks.size();
  }

  size_t totalTaskCount() const {
    Synchronized s(_monitor);   
    return _tasks.size() + _workerCount - _idleCount;
  }
  
  void add(shared_ptr<Runnable> value);

  void remove(shared_ptr<Runnable> task);

private:
  size_t _workerCount;
  size_t _workerMaxCount;
  size_t _idleCount;
  ThreadManager::STATE _state;
  shared_ptr<ThreadFactory> _threadFactory;


  friend class ThreadManager::Task;
  std::queue<shared_ptr<Task> > _tasks;
  Monitor _monitor;
  Monitor _workerMonitor;

  friend class ThreadManager::Worker;
  std::set<shared_ptr<Thread> > _workers;
  std::set<shared_ptr<Thread> > _deadWorkers;
};

class ThreadManager::Task : public Runnable {

 public:
  enum STATE {
    WAITING,
    EXECUTING,
    CANCELLED,
    COMPLETE
  };

  Task(shared_ptr<Runnable> runnable) :
    _runnable(runnable),
    _state(WAITING) {}

  ~Task() {}

  void run() {
    if (_state == EXECUTING) {
      _runnable->run();
      _state = COMPLETE;
    }
  }

 private:
  shared_ptr<Runnable> _runnable;
  friend class ThreadManager::Worker;
  STATE _state;
};

class ThreadManager::Worker: public Runnable {
  enum STATE {
    UNINITIALIZED,
    STARTING,
    STARTED,
    STOPPING,
    STOPPED
  };

 public:
  Worker(ThreadManager::Impl* manager) : 
    _manager(manager),
    _state(UNINITIALIZED),
    _idle(false) {}

  ~Worker() {}

  bool isActive() const { return _manager->_workerCount <= _manager->_workerMaxCount; }

  /**
   * Worker entry point
   *
   * As long as worker thread is running, pull tasks off the task queue and
   * execute.
   */
  void run() {
    bool active = false;  
    bool notifyManager = false;

    /**
     * Increment worker semaphore and notify manager if worker count reached
     * desired max
     *
     * Note: We have to release the monitor and acquire the workerMonitor
     * since that is what the manager blocks on for worker add/remove
     */
    {
      Synchronized s(_manager->_monitor);
      active = _manager->_workerCount < _manager->_workerMaxCount;
      if (active) {
	_manager->_workerCount++;
	notifyManager = _manager->_workerCount == _manager->_workerMaxCount;
      }
    }

    if (notifyManager) {
      Synchronized s(_manager->_workerMonitor);
      _manager->_workerMonitor.notify();
      notifyManager = false;
    }

    while (active) {
      shared_ptr<ThreadManager::Task> task;

      /**
       * While holding manager monitor block for non-empty task queue (Also
       * check that the thread hasn't been requested to stop). Once the queue
       * is non-empty, dequeue a task, release monitor, and execute. If the
       * worker max count has been decremented such that we exceed it, mark
       * ourself inactive, decrement the worker count and notify the manager
       * (technically we're notifying the next blocked thread but eventually
       * the manager will see it.
       */
      {
        Synchronized s(_manager->_monitor);
	active = isActive();
	while (active && _manager->_tasks.empty()) {
          _manager->_idleCount++;
	  _idle = true;
          _manager->_monitor.wait();
          active = isActive();
	  _idle = false;
          _manager->_idleCount--;
	}

	if (active) {
          if (!_manager->_tasks.empty()) {
            task = _manager->_tasks.front();
            _manager->_tasks.pop();
            if (task->_state == ThreadManager::Task::WAITING) {
              task->_state = ThreadManager::Task::EXECUTING;
	    }
	  }
	} else {
	  _idle = true;  
	  _manager->_workerCount--;
          notifyManager = _manager->_workerCount == _manager->_workerMaxCount;
	}
      }
      
      if (task != NULL) {
	if (task->_state == ThreadManager::Task::EXECUTING) {
	  try {
            task->run();
          } catch(...) {
            // XXX need to log this
	  }
	}
      }
    }
    
    {
      Synchronized s(_manager->_workerMonitor);    
      _manager->_deadWorkers.insert(this->thread());
      if (notifyManager) {
        _manager->_workerMonitor.notify();
      }
    }
    
    return;
  }
  
  private:
    ThreadManager::Impl* _manager;
    friend class ThreadManager::Impl;
    STATE _state;
    bool _idle;
};


  void ThreadManager::Impl::addWorker(size_t value) {
  std::set<shared_ptr<Thread> > newThreads;
  for (size_t ix = 0; ix < value; ix++) {
    class ThreadManager::Worker;     
    shared_ptr<ThreadManager::Worker> worker = shared_ptr<ThreadManager::Worker>(new ThreadManager::Worker(this));
    newThreads.insert(_threadFactory->newThread(worker));
  }

  {
    Synchronized s(_monitor);
    _workerMaxCount+= value;
    _workers.insert(newThreads.begin(), newThreads.end());
  }
  
  for (std::set<shared_ptr<Thread> >::iterator ix = newThreads.begin(); ix != newThreads.end(); ix++) {
    shared_ptr<ThreadManager::Worker> worker = dynamic_pointer_cast<ThreadManager::Worker, Runnable>((*ix)->runnable());
    worker->_state = ThreadManager::Worker::STARTING;
    (*ix)->start();
  }

  {
    Synchronized s(_workerMonitor); 
    while (_workerCount != _workerMaxCount) {
      _workerMonitor.wait();
    }
  }
}

void ThreadManager::Impl::start() {

  if (_state == ThreadManager::STOPPED) {
    return;
  }

  {
    Synchronized s(_monitor); 
    if (_state == ThreadManager::UNINITIALIZED) {
      if (_threadFactory == NULL) {
        throw InvalidArgumentException();
      }
      _state = ThreadManager::STARTED;
      _monitor.notifyAll();
    }

    while (_state == STARTING) {
      _monitor.wait();
    }
  }
}

void ThreadManager::Impl::stop() {
  bool doStop = false;
  if (_state == ThreadManager::STOPPED) {
    return;
  }

  {
    Synchronized s(_monitor); 
    if (!_state != ThreadManager::STOPPING && _state != ThreadManager::STOPPED) {
      doStop = true;
      _state = ThreadManager::STOPPING;
    }
  }

  if (doStop) {
    removeWorker(_workerCount);
    _state = ThreadManager::STOPPING;
  }

  // XXX 
  // should be able to block here for transition to STOPPED since we're no
  // using shared_ptrs
}
  
void ThreadManager::Impl::removeWorker(size_t value) {
  std::set<shared_ptr<Thread> > removedThreads;
  {
    Synchronized s(_monitor); 
    if (value > _workerMaxCount) {
      throw InvalidArgumentException();
    }

    _workerMaxCount-= value;

    if (_idleCount < value) {
      for (size_t ix = 0; ix < _idleCount; ix++) {
	_monitor.notify();
      }
    } else {
      _monitor.notifyAll();
    }
  }

  {
    Synchronized s(_workerMonitor); 

    while (_workerCount != _workerMaxCount) {
      _workerMonitor.wait();
    }

    for (std::set<shared_ptr<Thread> >::iterator ix = _deadWorkers.begin(); ix != _deadWorkers.end(); ix++) {
      _workers.erase(*ix);
    }
    
    _deadWorkers.clear();
  }
}
  
void ThreadManager::Impl::add(shared_ptr<Runnable> value) {
    Synchronized s(_monitor); 

    if (_state != ThreadManager::STARTED) {
      throw IllegalStateException();
    }

    _tasks.push(shared_ptr<ThreadManager::Task>(new ThreadManager::Task(value)));

    // If idle thread is available notify it, otherwise all worker threads are
    // running and will get around to this task in time.
    if (_idleCount > 0) {
      _monitor.notify();
    }
  }

void ThreadManager::Impl::remove(shared_ptr<Runnable> task) {
  Synchronized s(_monitor); 
  if (_state != ThreadManager::STARTED) {
    throw IllegalStateException();
  }
}

class SimpleThreadManager : public ThreadManager::Impl {

public:
  SimpleThreadManager(size_t workerCount=4) : 
    _workerCount(workerCount),
    _firstTime(true) {
  }

  void start() {
    ThreadManager::Impl::start();
    addWorker(_workerCount);
  }

private:
  const size_t _workerCount;
  bool _firstTime;
  Monitor _monitor;
};


shared_ptr<ThreadManager> ThreadManager::newThreadManager() {
  return shared_ptr<ThreadManager>(new ThreadManager::Impl());
}

shared_ptr<ThreadManager> ThreadManager::newSimpleThreadManager(size_t count) {
  return shared_ptr<ThreadManager>(new SimpleThreadManager(count));
}

}}} // facebook::thrift::concurrency