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"
#include "FWCore/Utilities/interface/thread_safety_macros.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)) {
          // Exception needs to be caught in order to decrease the waitTask reference count at the end. The user of SerialTaskQueue should handle exceptions within iAction.
          CMS_SA_ALLOW 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