nova-tt/thread_pool.hpp

   1 //  c++11-style thread pool
   2 //  Copyright (C) 2013 Tim Blechmann
   3 //
   4 //  This program is free software; you can redistribute it and/or modify
   5 //  it under the terms of the GNU General Public License as published by
   6 //  the Free Software Foundation; either version 2 of the License, or
   7 //  (at your option) any later version.
   8 //
   9 //  This program is distributed in the hope that it will be useful,
  10 //  but WITHOUT ANY WARRANTY; without even the implied warranty of
  11 //  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  12 //  GNU General Public License for more details.
  13 //
  14 //  You should have received a copy of the GNU General Public License
  15 //  along with this program; see the file COPYING.  If not, write to
  16 //  the Free Software Foundation, Inc., 59 Temple Place - Suite 330,
  17 //  Boston, MA 02111-1307, USA.
  18
  19 #ifndef THREAD_POOL_HPP
  20 #define THREAD_POOL_HPP
  21
  22 #include <condition_variable>
  23 #include <deque>
  24 #include <future>
  25 #include <mutex>
  26 #include <thread>
  27 #include <vector>
  28
  29 namespace nova   {
  30 namespace detail {
  31
  32 struct task_base
  33 {
  34     virtual void run() = 0;
  35 };
  36
  37 template <typename T>
  38 struct task:
  39     task_base
  40 {
  41     typedef std::packaged_task<T()> work;
  42
  43     explicit task (work && rhs):
  44         packaged_task(std::move(rhs))
  45     {}
  46
  47     void run()
  48     {
  49         packaged_task();
  50     }
  51
  52     std::packaged_task<T()> packaged_task;
  53 };
  54
  55 // LATER: implement some work-stealing
  56 class thread_pool
  57 {
  58 public:
  59     thread_pool(size_t thread_count):
  60         request_stop(false)
  61     {
  62         for (size_t i = 0; i != thread_count; ++i)
  63             worker_threads.emplace_back(std::bind(&thread_pool::worker_thread, this));
  64     }
  65
  66     ~thread_pool()
  67     {
  68         request_stop = true;
  69         condition.notify_all();
  70
  71         for (;;) {
  72             std::shared_ptr<task_base> work;
  73             {
  74                 std::unique_lock<std::mutex> lock(queue_guard);
  75                 if (tasks.empty())
  76                     break;
  77
  78                 work = std::move(tasks.front());
  79                 tasks.pop_front();
  80             }
  81             work->run();
  82         }
  83
  84         for (auto & thread : worker_threads)
  85             thread.join();
  86     }
  87
  88     template<class Functor>
  89     auto schedule(Functor f) -> std::future<decltype( f() )>
  90     {
  91         typedef decltype(f()) T;
  92
  93         if (worker_threads.empty())
  94             return std::async(std::launch::async, f);
  95
  96         std::packaged_task<T()> job(f);
  97         std::future<T> ret = job.get_future();
  98
  99         {
 100             std::unique_lock<std::mutex> lock(queue_guard);
 101             std::shared_ptr<task_base> task_ptr (new task<T>(std::move(job)));
 102             tasks.emplace_back(task_ptr);
 103         }
 104
 105         condition.notify_one();
 106         return ret;
 107     }
 108
 109 private:
 110     void worker_thread()
 111     {
 112         for (;;) {
 113             std::shared_ptr<task_base> work;
 114             {
 115                 std::unique_lock<std::mutex> lock(queue_guard);
 116                 while(!request_stop && tasks.empty())
 117                     condition.wait(lock);
 118
 119                 if (request_stop && tasks.empty())
 120                     return;
 121
 122                 work = std::move(tasks.front());
 123                 tasks.pop_front();
 124             }
 125             work->run();
 126         }
 127     }
 128
 129     std::vector< std::thread > worker_threads;
 130     std::deque< std::shared_ptr<task_base> > tasks;
 131
 132     std::mutex queue_guard;
 133     std::condition_variable condition;
 134     bool request_stop;
 135 };
 136
 137 }
 138
 139 using detail::thread_pool;
 140
 141 }
 142
 143 #endif // THREAD_POOL_HPP