thread_safe_queue.hpp

Go to the documentation of this file.

// Copyright (c) 2014-2023, Lawrence Livermore National Security, LLC.
// Produced at the Lawrence Livermore National Laboratory.
// Written by the LBANN Research Team (B. Van Essen, et al.) listed in
// the CONTRIBUTORS file. <lbann-dev@llnl.gov>
//
// LLNL-CODE-697807.
// All rights reserved.
//
// This file is part of LBANN: Livermore Big Artificial Neural Network
// Toolkit. For details, see http://software.llnl.gov/LBANN or
// https://github.com/LLNL/LBANN.
//
// Licensed under the Apache License, Version 2.0 (the "Licensee"); 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 LBANN_UTILS_THREADS_THREAD_SAFE_QUEUE_HPP_INCLUDED
#define LBANN_UTILS_THREADS_THREAD_SAFE_QUEUE_HPP_INCLUDED

#include <lbann/utils/memory.hpp>

#include <condition_variable>
#include <memory>
#include <mutex>

namespace lbann {

template <typename T>
class thread_safe_queue
{
private:
  struct _Node
  {
    std::unique_ptr<T> data_;
    std::unique_ptr<_Node> next_;
  };

public:
  thread_safe_queue()
    : head_(make_unique<_Node>()), tail_(head_.get()), m_stop_threads(false)
  {}

  void push(T value)
  {
    // Make the new data outside of the lock to minimize lock time
    auto new_value = std::make_unique<T>(std::move(value));
    auto new_node = std::make_unique<_Node>();

    // Adding to the queue only modifies the tail
    {
      std::lock_guard<std::mutex> lk(tail_mtx_);
      tail_->data_ = std::move(new_value);
      tail_->next_ = std::move(new_node);
      tail_ = tail_->next_.get();
    }
    // Update the condition variable (for wait_and_pop)
    data_available_.notify_one();
  }

  void wake_all(bool stop = false)
  {
    {
      std::lock_guard<std::mutex> lk(head_mtx_);
      m_stop_threads = stop;
    }
    // Update the condition variable (for wait_and_pop)
    data_available_.notify_all();
  }

  void set_stop_threads(bool flag) { m_stop_threads = flag; }

  std::unique_ptr<T> try_pop()
  {
    std::unique_lock<std::mutex> lk(head_mtx_);
    if (head_.get() == do_get_tail_())
      return nullptr;

    // Remove the head
    auto popped_head = std::move(head_);
    head_ = std::move(popped_head->next_);

    return std::move(popped_head->data_);
  }

  std::unique_ptr<T> wait_and_pop()
  {
    std::unique_lock<std::mutex> lk(head_mtx_);
    data_available_.wait(lk, [&] {
      return ((head_.get() != do_get_tail_()) ||
              ((head_.get() == do_get_tail_()) && m_stop_threads));
    });

    // There is no more work to do, bail
    if (head_.get() == do_get_tail_() && m_stop_threads) {
      return nullptr;
    }

    // Remove the head
    auto popped_head = std::move(head_);
    head_ = std::move(popped_head->next_);

    return std::move(popped_head->data_);
  }

  bool empty() const
  {
    std::lock_guard<std::mutex> lk(head_mtx_);
    return (head_.get() == do_get_tail_());
  }

private:
  _Node* do_get_tail_() const
  {
    std::lock_guard<std::mutex> lk(tail_mtx_);
    return tail_;
  }

private:
  mutable std::mutex head_mtx_;

  mutable std::mutex tail_mtx_;

  std::unique_ptr<_Node> head_;

  _Node* tail_;

  std::condition_variable data_available_;

  bool m_stop_threads;

}; // class thread_safe_queue

} // namespace lbann
#endif /* LBANN_UTILS_THREADS_THREAD_SAFE_QUEUE_HPP_INCLUDED */