LimitedTaskQueue.h

Go to the documentation of this file.

#ifndef FWCore_Concurrency_LimitedTaskQueue_h
#define FWCore_Concurrency_LimitedTaskQueue_h
// -*- C++ -*-
//
// Package:     Concurrency
// Class  :     LimitedTaskQueue
//
//
// Original Author:  Chris Jones
//         Created:  Thu Feb 21 11:14:39 CST 2013
// $Id$
//

// system include files
#include <atomic>
#include <vector>
#include <memory>

#include "FWCore/Concurrency/interface/SerialTaskQueue.h"

// user include files

// forward declarations
namespace edm {
  class LimitedTaskQueue {
  public:
    LimitedTaskQueue(unsigned int iLimit) : m_queues{iLimit} {}

    // ---------- member functions ---------------------------


    template <typename T>
    void push(T&& iAction);


    template <typename T>
    void pushAndWait(T&& iAction);

    class Resumer {
    public:
      friend class LimitedTaskQueue;

      Resumer() = default;
      ~Resumer() { resume(); }

      Resumer(Resumer&& iOther) : m_queue(iOther.m_queue) { iOther.m_queue = nullptr; }

      Resumer(Resumer const& iOther) : m_queue(iOther.m_queue) {
        if (m_queue) {
          m_queue->pause();
        }
      }

      Resumer& operator=(Resumer const& iOther) {
        auto t = iOther;
        return (*this = std::move(t));
      }
      Resumer& operator=(Resumer&& iOther) {
        if (m_queue) {
          m_queue->resume();
        }
        m_queue = iOther.m_queue;
        iOther.m_queue = nullptr;
        return *this;
      }

      bool resume() {
        if (m_queue) {
          auto q = m_queue;
          m_queue = nullptr;
          return q->resume();
        }
        return false;
      }

    private:
      Resumer(SerialTaskQueue* iQueue) : m_queue{iQueue} {}
      SerialTaskQueue* m_queue = nullptr;
    };


    template <typename T>
    void pushAndPause(T&& iAction);

    unsigned int concurrencyLimit() const { return m_queues.size(); }

  private:
    LimitedTaskQueue(const LimitedTaskQueue&) = delete;
    const LimitedTaskQueue& operator=(const LimitedTaskQueue&) = delete;

    // ---------- member data --------------------------------
    std::vector<SerialTaskQueue> m_queues;
  };

  template <typename T>
  void LimitedTaskQueue::push(T&& iAction) {
    auto set_to_run = std::make_shared<std::atomic<bool>>(false);
    for (auto& q : m_queues) {
      q.push([set_to_run, iAction]() mutable {
        bool expected = false;
        if (set_to_run->compare_exchange_strong(expected, true)) {
          iAction();
        }
      });
    }
  }

  template <typename T>
  void LimitedTaskQueue::pushAndWait(T&& iAction) {
    tbb::empty_task* waitTask = new (tbb::task::allocate_root()) tbb::empty_task;
    waitTask->set_ref_count(2);
    auto set_to_run = std::make_shared<std::atomic<bool>>(false);
    for (auto& q : m_queues) {
      q.push([set_to_run, waitTask, iAction]() mutable {
        bool expected = false;
        if (set_to_run->compare_exchange_strong(expected, true)) {
          try {
            iAction();
          } catch (...) {
          }
          waitTask->decrement_ref_count();
        }
      });
    }
    waitTask->wait_for_all();
    tbb::task::destroy(*waitTask);
  }

  template <typename T>
  void LimitedTaskQueue::pushAndPause(T&& iAction) {
    auto set_to_run = std::make_shared<std::atomic<bool>>(false);
    for (auto& q : m_queues) {
      q.push([&q, set_to_run, iAction]() mutable {
        bool expected = false;
        if (set_to_run->compare_exchange_strong(expected, true)) {
          q.pause();
          iAction(Resumer(&q));
        }
      });
    }
  }

}  // namespace edm

#endif