minrt/TaskDispatch.cpp at master · wheresthecode/minrt · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
// https://bitbucket.org/wolfpld/etcpak
// Bartosz Taudul <wolf.pld@gmail.com>

#include <assert.h>

#include "TaskDispatch.hpp"

static TaskDispatch* s_instance = nullptr;

TaskDispatch::TaskDispatch(size_t workers)
  : m_exit(false)
  , m_jobs(0)
{
  assert(!s_instance);
  s_instance = this;

  assert(workers >= 1);
  workers--;

  m_workers.reserve(workers);
  for (int i = 0; i < workers; i++)
  {
    m_workers.emplace_back([this]{ Worker(); });
  }
}

TaskDispatch::~TaskDispatch()
{
  m_exit = true;
  m_cvWork.notify_all();

  for (auto& worker : m_workers)
  {
    worker.join();
  }

  assert(s_instance);
  s_instance = nullptr;
}

void TaskDispatch::Queue(const std::function<void(void)>& f)
{
  std::unique_lock<std::mutex> lock(s_instance->m_queueLock);
  s_instance->m_queue.emplace(f);
  const auto size = s_instance->m_queue.size();
  lock.unlock();
  if (size > 1)
  {
    s_instance->m_cvWork.notify_one();
  }
}

void TaskDispatch::Queue(std::function<void(void)>&& f)
{
  std::unique_lock<std::mutex> lock(s_instance->m_queueLock);
  s_instance->m_queue.emplace(std::move(f));
  const auto size = s_instance->m_queue.size();
  lock.unlock();
  if (size > 1)
  {
    s_instance->m_cvWork.notify_one();
  }
}

void TaskDispatch::Sync()
{
  std::unique_lock<std::mutex> lock(s_instance->m_queueLock);
  while (!s_instance->m_queue.empty())
  {
    auto f = s_instance->m_queue.front();
    s_instance->m_queue.pop();
    lock.unlock();
    f();
    lock.lock();
  }
  s_instance->m_cvJobs.wait(lock, []{ return s_instance->m_jobs == 0; });
}

void TaskDispatch::Worker()
{
  for (;;)
  {
    std::unique_lock<std::mutex> lock(m_queueLock);
    m_cvWork.wait(lock, [this]{ return !m_queue.empty() || m_exit; });
    if (m_exit) return;
    auto f = m_queue.front();
    m_queue.pop();
    m_jobs++;
    lock.unlock();
    f();
    lock.lock();
    m_jobs--;
    bool notify = m_jobs == 0 && m_queue.empty();
    lock.unlock();
    if (notify)
    {
      m_cvJobs.notify_all();
    }
  }
}