thread_group.h

Go to the documentation of this file.

#ifndef DMLC_THREAD_GROUP_H_
#define DMLC_THREAD_GROUP_H_

#include <dmlc/concurrentqueue.h>
#include <dmlc/blockingconcurrentqueue.h>
#include <dmlc/logging.h>
#include <string>
#include <mutex>
#include <utility>
#include <memory>
#include <set>
#include <thread>
#include <unordered_set>
#include <unordered_map>
#if defined(DMLC_USE_CXX14) || __cplusplus > 201103L  /* C++14 */
#include <shared_mutex>
#endif
#include <condition_variable>
#ifdef __linux__
#include <unistd.h>
#include <sys/syscall.h>
#endif

namespace dmlc {

class ManualEvent {
 public:
  ManualEvent() : signaled_(false) {}

  void wait() {
    std::unique_lock<std::mutex> lock(mutex_);
    if (!signaled_) {
      condition_variable_.wait(lock);
    }
  }

  void signal() {
    signaled_ = true;
    std::unique_lock<std::mutex> lk(mutex_);
    condition_variable_.notify_all();
  }

  void reset() {
    std::unique_lock<std::mutex> lk(mutex_);
    signaled_ = false;
  }

 private:
  std::mutex mutex_;
  std::condition_variable condition_variable_;
  std::atomic<bool> signaled_;
};

#if defined(DMLC_USE_CXX14) || __cplusplus > 201103L  /* C++14 */

using SharedMutex = std::shared_timed_mutex;
using WriteLock = std::unique_lock<SharedMutex>;
using ReadLock = std::shared_lock<SharedMutex>;
#else

using SharedMutex = std::recursive_mutex;
using WriteLock = std::unique_lock<SharedMutex>;
using ReadLock = std::unique_lock<SharedMutex>;
#endif

class ThreadGroup {
 public:
  class Thread {
   public:
    using SharedPtr = std::shared_ptr<Thread>;

    Thread(std::string threadName, ThreadGroup *owner, std::thread *thrd = nullptr)
      : name_(std::move(threadName))
        , thread_(thrd)
        , ready_event_(std::make_shared<ManualEvent>())
        , start_event_(std::make_shared<ManualEvent>())
        , owner_(owner)
        , shutdown_requested_(false)
        , auto_remove_(false) {
      CHECK_NOTNULL(owner);
    }

    virtual ~Thread() {
      const bool self_delete = is_current_thread();
      if (!self_delete) {
        request_shutdown();
        internal_join(true);
      }
      WriteLock guard(thread_mutex_);
      if (thread_.load()) {
        std::thread *thrd = thread_.load();
        thread_ = nullptr;
        if (self_delete) {
          thrd->detach();
        }
        delete thrd;
      }
    }

    const char *name() const {
      return name_.c_str();
    }

    template<typename StartFunction, typename ...Args>
    static bool launch(std::shared_ptr<Thread> pThis,
                       bool autoRemove,
                       StartFunction start_function,
                       Args ...args);

    bool is_current_thread() const {
      ReadLock guard(thread_mutex_);
      return thread_.load() ? (thread_.load()->get_id() == std::this_thread::get_id()) : false;
    }

    virtual void request_shutdown() {
      shutdown_requested_ = true;
    }

    virtual bool is_shutdown_requested() const {
      return shutdown_requested_.load();
    }

    bool is_auto_remove() const {
      return auto_remove_;
    }

    void make_joinable() {
      auto_remove_ = false;
    }

    bool joinable() const {
      if (thread_.load()) {
        CHECK_EQ(auto_remove_, false);
        // be checked by searching the group or exit event.
        return thread_.load()->joinable();
      }
      return false;
    }

    void join() {
      internal_join(false);
    }

    std::thread::id get_id() const {
      return thread_.load()->get_id();
    }

   private:
    void internal_join(bool auto_remove_ok) {
      ReadLock guard(thread_mutex_);
      // should be careful calling (or any function externally) this when in
      // auto-remove mode
      if (thread_.load() && thread_.load()->get_id() != std::thread::id()) {
        std::thread::id someId;
        if (!auto_remove_ok) {
          CHECK_EQ(auto_remove_, false);
        }
        CHECK_NOTNULL(thread_.load());
        if (thread_.load()->joinable()) {
          thread_.load()->join();
        } else {
          LOG(WARNING) << "Thread " << name_ << " ( "
                       << thread_.load()->get_id() << " ) not joinable";
        }
      }
    }

    template <typename StartFunction, typename ...Args>
    static int entry_and_exit_f(std::shared_ptr<Thread> pThis,
                                StartFunction start_function,
                                Args... args);
    std::string name_;
    mutable SharedMutex thread_mutex_;
    std::atomic<std::thread *> thread_;
    std::shared_ptr<ManualEvent> ready_event_;
    std::shared_ptr<ManualEvent> start_event_;
    ThreadGroup *owner_;
    std::atomic<bool> shutdown_requested_;
    std::atomic<bool> auto_remove_;
  };

  inline ThreadGroup()
    : evEmpty_(std::make_shared<ManualEvent>()) {
    evEmpty_->signal();  // Starts out empty
  }

  virtual ~ThreadGroup() {
    request_shutdown_all();
    join_all();
  }

  inline bool is_this_thread_in() const {
    std::thread::id id = std::this_thread::get_id();
    ReadLock guard(m_);
    for (auto it = threads_.begin(), end = threads_.end(); it != end; ++it) {
      std::shared_ptr<Thread> thrd = *it;
      if (thrd->get_id() == id)
        return true;
    }
    return false;
  }

  inline bool is_thread_in(std::shared_ptr<Thread> thrd) const {
    if (thrd) {
      std::thread::id id = thrd->get_id();
      ReadLock guard(m_);
      for (auto it = threads_.begin(), end = threads_.end(); it != end; ++it) {
        std::shared_ptr<Thread> thrd = *it;
        if (thrd->get_id() == id)
          return true;
      }
      return false;
    } else {
      return false;
    }
  }

  inline bool add_thread(std::shared_ptr<Thread> thrd) {
    if (thrd) {
      WriteLock guard(m_);
      auto iter = name_to_thread_.find(thrd->name());
      if (iter == name_to_thread_.end()) {
        name_to_thread_.emplace(std::make_pair(thrd->name(), thrd));
        CHECK_EQ(threads_.insert(thrd).second, true);
        evEmpty_->reset();
        return true;
      }
    }
    return false;
  }

  inline bool remove_thread(std::shared_ptr<Thread> thrd) {
    if (thrd) {
      WriteLock guard(m_);
      auto iter = threads_.find(thrd);
      if (iter != threads_.end()) {
        name_to_thread_.erase(thrd->name());
        threads_.erase(iter);
        if (threads_.empty()) {
          evEmpty_->signal();
        }
        return true;
      }
    }
    return false;
  }

  inline void join_all() {
    CHECK_EQ(!is_this_thread_in(), true);
    do {
      std::unique_lock<std::mutex> lk(join_all_mtx_);
      std::unordered_set<std::shared_ptr<Thread>> working_set;
      {
        ReadLock guard(m_);
        for (auto iter = threads_.begin(), e_iter = threads_.end(); iter != e_iter; ++iter) {
          if (!(*iter)->is_auto_remove()) {
            working_set.emplace(*iter);
          }
        }
      }
      // Where possible, prefer to do a proper join rather than simply waiting for empty
      // (easier to troubleshoot)
      while (!working_set.empty()) {
        std::shared_ptr<Thread> thrd;
        thrd = *working_set.begin();
        if (thrd->joinable()) {
          thrd->join();
        }
        remove_thread(thrd);
        working_set.erase(working_set.begin());
        thrd.reset();
      }
      // Wait for auto-remove threads (if any) to complete
    } while (0);
    evEmpty_->wait();
    CHECK_EQ(threads_.size(), 0);
  }

  inline void request_shutdown_all(const bool make_all_joinable = true) {
    std::unique_lock<std::mutex> lk(join_all_mtx_);
    ReadLock guard(m_);
    for (auto &thread : threads_) {
      if (make_all_joinable) {
        thread->make_joinable();
      }
      thread->request_shutdown();
    }
  }

  inline size_t size() const {
    ReadLock guard(m_);
    return threads_.size();
  }

  inline bool empty() const {
    ReadLock guard(m_);
    return threads_.size() == 0;
  }

  template<typename StartFunction, typename ThreadType = Thread, typename ...Args>
  inline bool create(const std::string &threadName,
                     bool auto_remove,
                     StartFunction start_function,
                     Args... args) {
    typename ThreadType::SharedPtr newThread(new ThreadType(threadName, this));
    return Thread::launch(newThread, auto_remove, start_function, args...);
  }

  inline std::shared_ptr<Thread> thread_by_name(const std::string& name) {
    ReadLock guard(m_);
    auto iter = name_to_thread_.find(name);
    if (iter != name_to_thread_.end()) {
      return iter->second;
    }
    return nullptr;
  }

 private:
  mutable SharedMutex m_;
  mutable std::mutex join_all_mtx_;
  std::unordered_set<std::shared_ptr<Thread>> threads_;
  std::shared_ptr<ManualEvent> evEmpty_;
  std::unordered_map<std::string, std::shared_ptr<Thread>> name_to_thread_;
};

template<typename ObjectType, ObjectType quit_item>
class BlockingQueueThread : public ThreadGroup::Thread {
  using BQT = BlockingQueueThread<ObjectType, quit_item>;

 public:
  BlockingQueueThread(const std::string& name,
                      dmlc::ThreadGroup *owner,
                      std::thread *thrd = nullptr)
    : ThreadGroup::Thread(std::move(name), owner, thrd)
      , shutdown_in_progress_(false) {
  }


  ~BlockingQueueThread() override {
    // Call to parent first because we don't want to wait for the queue to empty
    ThreadGroup::Thread::request_shutdown();
    request_shutdown();
  }

  void request_shutdown() override {
    shutdown_in_progress_ = true;
    while (queue_->size_approx() > 0 && !ThreadGroup::Thread::is_shutdown_requested()) {
      std::this_thread::sleep_for(std::chrono::milliseconds(1));
    }
    ThreadGroup::Thread::request_shutdown();
    queue_->enqueue(quit_item);
  }

  void enqueue(const ObjectType& item) {
    if (!shutdown_in_progress_) {
      queue_->enqueue(item);
    }
  }

  size_t size_approx() const { return queue_->size_approx(); }

  template<typename SecondaryFunction>
  static bool launch_run(std::shared_ptr<BQT> pThis,
                         SecondaryFunction secondary_function) {
    return ThreadGroup::Thread::launch(pThis, true, [](std::shared_ptr<BQT> pThis,
                                                       SecondaryFunction secondary_function) {
                                         return pThis->run(secondary_function);
                                       },
                                       pThis, secondary_function);
  }

  template<typename OnItemFunction>
  inline int run(OnItemFunction on_item_function) {
    int rc = 0;
    do {
      ObjectType item;
      queue_->wait_dequeue(item);
      if (item == quit_item) {
        break;
      }
      rc = on_item_function(item);
      if (rc) {
        break;
      }
    } while (true);
    return rc;
  }

 private:
  std::shared_ptr<dmlc::moodycamel::BlockingConcurrentQueue<ObjectType>> queue_ =
    std::make_shared<dmlc::moodycamel::BlockingConcurrentQueue<ObjectType>>();
  std::atomic<bool> shutdown_in_progress_;
};

template<typename Duration>
class TimerThread : public ThreadGroup::Thread {
  using ThreadGroup::Thread::is_shutdown_requested;

 public:
  TimerThread(const std::string& name, ThreadGroup *owner)
    : Thread(name, owner) {
  }

  ~TimerThread() override {
    request_shutdown();
  }

  template<typename SecondaryFunction>
  static bool launch_run(std::shared_ptr<TimerThread<Duration>> pThis,
                         SecondaryFunction secondary_function) {
    return ThreadGroup::Thread::launch(pThis, true, [](std::shared_ptr<TimerThread<Duration>> pThis,
                                                       SecondaryFunction secondary_function) {
                                         return pThis->run(secondary_function);
                                       },
                                       pThis, secondary_function);
  }

  template<typename Function>
  static void start(std::shared_ptr<TimerThread> timer_thread,
                    Duration duration,
                    Function function) {
    timer_thread->duration_ = duration;
    launch_run(timer_thread, function);
  }

  template<typename OnTimerFunction>
  inline int run(OnTimerFunction on_timer_function) {
    int rc = 0;
    while (!is_shutdown_requested()) {
      std::this_thread::sleep_for(duration_);
      if (!is_shutdown_requested()) {
        rc = on_timer_function();
      }
    }
    return rc;
  }

 private:
  Duration duration_;
};

/*
 * Inline functions - see declarations for usage
 */
template <typename StartFunction, typename ...Args>
inline int ThreadGroup::Thread::entry_and_exit_f(std::shared_ptr<Thread> pThis,
                                                 StartFunction start_function,
                                                 Args... args) {
  int rc;
  if (pThis) {
    // Signal launcher that we're up and running
    pThis->ready_event_->signal();
    // Wait for launcher to be ready for us to start
    pThis->start_event_->wait();
    // Reset start_event_ for possible reuse
    pThis->start_event_->reset();  // Reset in case it needs to be reused
    // If we haven't been requested to shut down prematurely, then run the desired function
    if (!pThis->is_shutdown_requested()) {
      rc = start_function(args...);
    } else {
      rc = -1;
    }
    // If we're set up as auto-remove, then remove this thread from the thread group
    if (pThis->is_auto_remove()) {
      pThis->owner_->remove_thread(pThis);
    }
    // Release this thread shared pinter. May or may not be the last reference.
    pThis.reset();
  } else {
    LOG(ERROR) << "Null pThis thread pointer";
    rc = EINVAL;
  }
  return rc;
}

template<typename StartFunction, typename ...Args>
inline bool ThreadGroup::Thread::launch(std::shared_ptr<Thread> pThis,
                                        bool autoRemove,
                                        StartFunction start_function,
                                        Args ...args) {
  WriteLock guard(pThis->thread_mutex_);
  CHECK_EQ(!pThis->thread_.load(), true);
  CHECK_NOTNULL(pThis->owner_);
  // Set auto remove
  pThis->auto_remove_ = autoRemove;
  // Create the actual stl thread object
  pThis->thread_ = new std::thread(Thread::template entry_and_exit_f<
                                     StartFunction, Args...>,
                                   pThis,
                                   start_function,
                                   args...);
  // Attempt to add the thread to the thread group (after started, since in case
  // something goes wrong, there's not a zombie thread in the thread group)
  if (!pThis->owner_->add_thread(pThis)) {
    pThis->request_shutdown();
    LOG(ERROR) << "Duplicate thread name within the same thread group is not allowed";
  }
  // Wait for the thread to spin up
  pThis->ready_event_->wait();
  // Signal the thgread to continue (it will check its shutdown status)
  pThis->start_event_->signal();
  // Return if successful
  return pThis->thread_.load() != nullptr;
}

template<typename Duration, typename TimerFunction>
inline bool CreateTimer(const std::string& timer_name,
                        const Duration& duration,
                        ThreadGroup *owner,
                        TimerFunction timer_function) {
  std::shared_ptr<dmlc::TimerThread<Duration>> timer_thread =
    std::make_shared<dmlc::TimerThread<Duration>>(timer_name, owner);
  dmlc::TimerThread<Duration>::start(timer_thread, duration, timer_function);
  return timer_thread != nullptr;
}
}  // namespace dmlc

#endif  // DMLC_THREAD_GROUP_H_