workerpool.cc

Go to the documentation of this file.

#include "ion/base/workerpool.h"

#include "ion/base/lockguards.h"

namespace ion {
namespace base {

WorkerPool::WorkerPool(WorkerPool::Worker* worker)
    : worker_(CHECK_NOTNULL(worker)),
      threads_(GetNonNullAllocator()),
      suspended_(true),
      killing_(false),
      slow_path_(false),
      spawn_func_([this](){ this->ThreadEntryPoint(); return true; }) {}

WorkerPool::~WorkerPool() {
  ion::base::LockGuard lock(&mutex_);
  KillAllThreads();
}

void WorkerPool::ResizeThreadPool(size_t thread_count) {
  ion::base::LockGuard lock(&mutex_);

  if (thread_count < threads_.size()) {
    KillAllThreads();
  }

  while (thread_count > threads_.size()) {
    threads_.insert(port::SpawnThreadStd(&spawn_func_));
    if (!suspended_)
      Post(&active_threads_sema_);
  }
}

void WorkerPool::Suspend() {
  ion::base::LockGuard lock(&mutex_);
  if (!suspended_) {
    suspended_ = true;
    slow_path_ = true;
    for (size_t i = 0; i < threads_.size(); ++i) {
      Wait(&active_threads_sema_);
    }
    slow_path_ = false;
  }
}

void WorkerPool::Resume() {
  ion::base::LockGuard lock(&mutex_);
  if (suspended_) {
    suspended_ = false;
    for (size_t i = 0; i < threads_.size(); ++i) {
      Post(&active_threads_sema_);
    }
  }
}

bool WorkerPool::IsSuspended() const {
  ion::base::LockGuard lock(&mutex_);
  return suspended_;
}

void WorkerPool::ThreadEntryPoint() {
  if (port::IsThreadNamingSupported())
    port::SetThreadName(GetName());

  while (true) {
    while (slow_path_) {
      if (killing_) {
        return;
      } else {
        port::YieldThread();
      }
    }

    Wait(&work_sema_);
    Wait(&active_threads_sema_);
    worker_->DoWork();
    Post(&active_threads_sema_);
  }
}

void WorkerPool::KillAllThreads() {
  DCHECK(mutex_.IsLocked());

  killing_ = true;
  slow_path_ = true;
  for (size_t i = 0; i < threads_.size(); ++i) {
    Post(&work_sema_);
    Post(&active_threads_sema_);
  }

  for (const auto& thread_id : threads_) {
    bool join_succeeded = port::JoinThread(thread_id);
    DCHECK(join_succeeded);
  }
  threads_.clear();
  slow_path_ = false;
  killing_ = false;

  while (active_threads_sema_.TryWait()) {}
}

void WorkerPool::Wait(port::Semaphore* sema) {
  bool semaphore_wait_succeeded = sema->Wait();
  DCHECK(semaphore_wait_succeeded);
}

void WorkerPool::Post(port::Semaphore* sema) {
  bool semaphore_post_succeeded = sema->Post();
  DCHECK(semaphore_post_succeeded);
}

}  // namespace base
}  // namespace ion