threadediter.h

Go to the documentation of this file.

#ifndef DMLC_THREADEDITER_H_
#define DMLC_THREADEDITER_H_
// defines DMLC_USE_CXX11
#include "./base.h"
// this code depends on c++11
#if DMLC_ENABLE_STD_THREAD
#include <condition_variable>
#include <functional>
#include <mutex>
#include <queue>
#include <atomic>
#include <thread>
#include <utility>
#include <memory>
#include "./data.h"
#include "./logging.h"

namespace dmlc {

class ScopedThread {
 public:
  explicit ScopedThread(std::thread thread)
      : thread_(std::move(thread)) {
    if (!thread_.joinable()) {
      throw std::logic_error("No thread");
    }
  }
  // destructor: join upon destruction
  virtual ~ScopedThread() {
    thread_.join();
  }
  // copy assignment and construction are not allowed
  ScopedThread(ScopedThread const&) = delete;
  ScopedThread& operator=(ScopedThread const&) = delete;

 private:
  std::thread thread_;
};

template<typename DType>
class ThreadedIter : public DataIter<DType> {
 public:
  class Producer {
   public:
    // virtual destructor
    virtual ~Producer() = default;
    virtual void BeforeFirst(void) {
      NotImplemented();
    }
    virtual bool Next(DType **inout_dptr) = 0;
  };
  explicit ThreadedIter(size_t max_capacity = 8)
      : producer_(nullptr),
        producer_thread_(nullptr),
        max_capacity_(max_capacity),
        nwait_consumer_(0),
        nwait_producer_(0),
        out_data_(NULL) {}
  virtual ~ThreadedIter(void) {
    this->Destroy();
  }
  inline void Destroy(void);
  inline void set_max_capacity(size_t max_capacity) {
    max_capacity_ = max_capacity;
  }
  inline void Init(std::shared_ptr<Producer> producer);
  inline void Init(std::function<bool(DType **)> next,
                   std::function<void()> beforefirst = NotImplemented);
  inline bool Next(DType **out_dptr);
  inline void Recycle(DType **inout_dptr);

  inline void ThrowExceptionIfSet(void);

  inline void ClearException(void);

  virtual bool Next(void) {
    if (out_data_ != NULL) {
      this->Recycle(&out_data_);
    }
    if (Next(&out_data_)) {
      return true;
    } else {
      return false;
    }
  }
  virtual const DType &Value(void) const {
    CHECK(out_data_ != NULL) << "Calling Value at beginning or end?";
    return *out_data_;
  }
  virtual void BeforeFirst(void) {
    ThrowExceptionIfSet();
    std::unique_lock<std::mutex> lock(mutex_);
    if (out_data_ != NULL) {
      free_cells_.push(out_data_);
      out_data_ = NULL;
    }
    if (producer_sig_.load(std::memory_order_acquire) == kDestroy)  return;

    producer_sig_.store(kBeforeFirst, std::memory_order_release);
    CHECK(!producer_sig_processed_.load(std::memory_order_acquire));
    if (nwait_producer_ != 0) {
      producer_cond_.notify_one();
    }
    CHECK(!producer_sig_processed_.load(std::memory_order_acquire));
    // wait until the request has been processed
    consumer_cond_.wait(lock, [this]() {
        return producer_sig_processed_.load(std::memory_order_acquire);
      });
    producer_sig_processed_.store(false, std::memory_order_release);
    bool notify = nwait_producer_ != 0 && !produce_end_;
    lock.unlock();
    // notify producer, in case they are waiting for the condition.
    if (notify) producer_cond_.notify_one();
    ThrowExceptionIfSet();
  }

 private:
  inline static void NotImplemented(void) {
    LOG(FATAL) << "BeforeFirst is not supported";
  }
  enum Signal {
    kProduce,
    kBeforeFirst,
    kDestroy
  };
  // Producer *producer_owned_;
  std::shared_ptr<Producer> producer_;

  std::atomic<Signal> producer_sig_;
  std::atomic<bool> producer_sig_processed_;
  std::unique_ptr<ScopedThread> producer_thread_;
  std::atomic<bool> produce_end_;
  size_t max_capacity_;
  std::mutex mutex_;
  std::mutex mutex_exception_;
  unsigned nwait_consumer_;
  unsigned nwait_producer_;
  std::condition_variable producer_cond_;
  std::condition_variable consumer_cond_;
  DType *out_data_;
  std::queue<DType*> queue_;
  std::queue<DType*> free_cells_;
  std::exception_ptr iter_exception_{nullptr};
};

// implementation of functions
template <typename DType> inline void ThreadedIter<DType>::Destroy(void) {
  if (producer_thread_) {
    {
      // lock the mutex
      std::lock_guard<std::mutex> lock(mutex_);
      // send destroy signal
      producer_sig_.store(kDestroy, std::memory_order_release);
      if (nwait_producer_ != 0) {
        producer_cond_.notify_one();
      }
    }
    producer_thread_.reset(nullptr);
  }
  // end of critical region
  // now the slave thread should exit
  while (free_cells_.size() != 0) {
    delete free_cells_.front();
    free_cells_.pop();
  }
  while (queue_.size() != 0) {
    delete queue_.front();
    queue_.pop();
  }
  if (producer_ != NULL) {
    producer_.reset();
  }
  if (out_data_ != NULL) {
    delete out_data_;
    out_data_ = NULL;
  }
}

template<typename DType>
inline void ThreadedIter<DType>::
Init(std::shared_ptr<Producer> producer) {
  CHECK(producer_ == NULL) << "can only call Init once";
  auto next = [producer](DType **dptr) {
      return producer->Next(dptr);
  };
  auto beforefirst = [producer]() {
    producer->BeforeFirst();
  };
  this->Init(next, beforefirst);
}

template <typename DType>
inline void ThreadedIter<DType>::Init(std::function<bool(DType **)> next,
                                      std::function<void()> beforefirst) {
  producer_sig_.store(kProduce, std::memory_order_release);
  producer_sig_processed_.store(false, std::memory_order_release);
  produce_end_.store(false, std::memory_order_release);
  ClearException();
  // procedure running in prodcuer
  // run producer thread
  auto producer_fun = [this, next, beforefirst]() {
    while (true) {
      try {
        DType *cell = NULL;
        {
          // lockscope
          std::unique_lock<std::mutex> lock(mutex_);
          ++this->nwait_producer_;
          producer_cond_.wait(lock, [this]() {
            if (producer_sig_.load(std::memory_order_acquire) == kProduce) {
              bool ret = !produce_end_.load(std::memory_order_acquire)
                         && (queue_.size() < max_capacity_ ||
                             free_cells_.size() != 0);
              return ret;
            } else {
              return true;
            }
          });
          --this->nwait_producer_;
          if (producer_sig_.load(std::memory_order_acquire) == kProduce) {
            if (free_cells_.size() != 0) {
              cell = free_cells_.front();
              free_cells_.pop();
            }
          } else if (producer_sig_.load(std::memory_order_acquire) == kBeforeFirst) {
            // reset the producer
            beforefirst();
            // cleanup the queue
            while (queue_.size() != 0) {
              free_cells_.push(queue_.front());
              queue_.pop();
            }
            // reset the state
            produce_end_.store(false, std::memory_order_release);
            producer_sig_processed_.store(true, std::memory_order_release);
            producer_sig_.store(kProduce, std::memory_order_release);
            // notify consumer that all the process as been done.
            lock.unlock();
            consumer_cond_.notify_all();
            continue;
          } else {
            // destroy the thread
            DCHECK(producer_sig_.load(std::memory_order_acquire) == kDestroy);
            producer_sig_processed_.store(true, std::memory_order_release);
            produce_end_.store(true, std::memory_order_release);
            lock.unlock();
            consumer_cond_.notify_all();
            return;
          }
        }  // end of lock scope
        // now without lock
        produce_end_.store(!next(&cell), std::memory_order_release);
        DCHECK(cell != NULL || produce_end_.load(std::memory_order_acquire));
        bool notify;
        {
          // lockscope
          std::lock_guard<std::mutex> lock(mutex_);
          if (!produce_end_.load(std::memory_order_acquire)) {
            queue_.push(cell);
          } else {
            if (cell != NULL)
              free_cells_.push(cell);
          }
          // put things into queue
          notify = nwait_consumer_ != 0;
        }
        if (notify)
          consumer_cond_.notify_all();
      } catch (std::exception &e) {
        // Shouldn't throw exception in destructor
        DCHECK(producer_sig_.load(std::memory_order_acquire) != kDestroy);
        {
          std::lock_guard<std::mutex> lock(mutex_exception_);
          if (!iter_exception_) {
            iter_exception_ = std::current_exception();
          }
        }
        bool next_notify = false;
        {
          std::unique_lock<std::mutex> lock(mutex_);
          if (producer_sig_.load(std::memory_order_acquire) == kBeforeFirst) {
            while (queue_.size() != 0) {
              free_cells_.push(queue_.front());
              queue_.pop();
            }
            produce_end_.store(true, std::memory_order_release);
            producer_sig_processed_.store(true, std::memory_order_release);
            lock.unlock();
            consumer_cond_.notify_all();
          } else if (producer_sig_.load(std::memory_order_acquire) == kProduce) {
            produce_end_.store(true, std::memory_order_release);
            next_notify = nwait_consumer_ != 0;
            lock.unlock();
            if (next_notify)
              consumer_cond_.notify_all();
          }
        }
        return;
      }
    }
  };
  producer_thread_.reset(new ScopedThread{std::thread(producer_fun)});
}

template <typename DType>
inline bool ThreadedIter<DType>::Next(DType **out_dptr) {
  if (producer_sig_.load(std::memory_order_acquire) == kDestroy)
    return false;
  ThrowExceptionIfSet();
  std::unique_lock<std::mutex> lock(mutex_);
  CHECK(producer_sig_.load(std::memory_order_acquire) == kProduce)
      << "Make sure you call BeforeFirst not inconcurrent with Next!";
  ++nwait_consumer_;
  consumer_cond_.wait(lock,
                      [this]() { return queue_.size() != 0
                                 || produce_end_.load(std::memory_order_acquire); });
  --nwait_consumer_;
  if (queue_.size() != 0) {
    *out_dptr = queue_.front();
    queue_.pop();
    bool notify = nwait_producer_ != 0
                  && !produce_end_.load(std::memory_order_acquire);
    lock.unlock();
    if (notify)
      producer_cond_.notify_one();

    ThrowExceptionIfSet();
    return true;
  } else {
    CHECK(produce_end_.load(std::memory_order_acquire));
    lock.unlock();

    ThrowExceptionIfSet();
    return false;
  }
}

template <typename DType>
inline void ThreadedIter<DType>::Recycle(DType **inout_dptr) {
  bool notify;
  ThrowExceptionIfSet();
  {
    std::lock_guard<std::mutex> lock(mutex_);
    free_cells_.push(*inout_dptr);
    *inout_dptr = NULL;
    notify = nwait_producer_ != 0 && !produce_end_.load(std::memory_order_acquire);
  }
  if (notify)
    producer_cond_.notify_one();
  ThrowExceptionIfSet();
}

template <typename DType> inline void ThreadedIter<DType>::ThrowExceptionIfSet(void) {
  std::exception_ptr tmp_exception{nullptr};
  {
    std::lock_guard<std::mutex> lock(mutex_exception_);
    if (iter_exception_) {
      tmp_exception = iter_exception_;
    }
  }
  if (tmp_exception) {
    try {
      std::rethrow_exception(tmp_exception);
    } catch (std::exception& exc) {
      LOG(FATAL) << exc.what();
    }
  }
}

template <typename DType> inline void ThreadedIter<DType>::ClearException(void) {
  std::lock_guard<std::mutex> lock(mutex_exception_);
  iter_exception_ = nullptr;
}

}  // namespace dmlc
#endif  // DMLC_USE_CXX11
#endif  // DMLC_THREADEDITER_H_