thread-pool.h

Go to the documentation of this file.

// Copyright 2012 Cloudera Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef IMPALA_UTIL_THREAD_POOL_H
#define IMPALA_UTIL_THREAD_POOL_H

#include "util/blocking-queue.h"

#include <boost/thread/mutex.hpp>
#include <boost/bind/mem_fn.hpp>

#include "util/thread.h"

namespace impala {

template <typename T>
class ThreadPool {
 public:
  typedef boost::function<void (int thread_id, const T& workitem)> WorkFunction;

  ThreadPool(const std::string& group, const std::string& thread_prefix,
      uint32_t num_threads, uint32_t queue_size, const WorkFunction& work_function)
    : work_function_(work_function),
      work_queue_(queue_size),
      shutdown_(false) {
    for (int i = 0; i < num_threads; ++i) {
      std::stringstream threadname;
      threadname << thread_prefix << "(" << i + 1 << ":" << num_threads << ")";
      threads_.AddThread(new Thread(group, threadname.str(),
          boost::bind<void>(boost::mem_fn(&ThreadPool<T>::WorkerThread), this, i)));
    }
  }

  ~ThreadPool() {
    Shutdown();
    Join();
  }

  //
  //
  bool Offer(const T& work) {
    return work_queue_.BlockingPut(work);
  }

  void Shutdown() {
    {
      boost::lock_guard<boost::mutex> l(lock_);
      shutdown_ = true;
    }
    work_queue_.Shutdown();
  }

  void Join() {
    threads_.JoinAll();
  }

  uint32_t GetQueueSize() const {
    return work_queue_.GetSize();
  }

  void DrainAndShutdown() {
    {
      boost::unique_lock<boost::mutex> l(lock_);
      while (work_queue_.GetSize() != 0) {
        empty_cv_.wait(l);
      }
    }
    Shutdown();
    Join();
  }

  Status AssignToCgroup(const std::string& cgroup) {
    return threads_.SetCgroup(cgroup);
  }

 private:
  void WorkerThread(int thread_id) {
    while (!IsShutdown()) {
      T workitem;
      if (work_queue_.BlockingGet(&workitem)) {
        work_function_(thread_id, workitem);
      }
      if (work_queue_.GetSize() == 0) {
        boost::unique_lock<boost::mutex> l(lock_);
        empty_cv_.notify_all();
      }
    }
  }

  bool IsShutdown() {
    boost::lock_guard<boost::mutex> l(lock_);
    return shutdown_;
  }

  WorkFunction work_function_;

  BlockingQueue<T> work_queue_;

  ThreadGroup threads_;

  boost::mutex lock_;

  bool shutdown_;

  boost::condition_variable empty_cv_;
};

}

#endif