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
101
102
103
104
105
106
107
108
109
110
111
112
113
114
/**
* Copyright (c) 2015-present, Facebook, Inc.
* All rights reserved.
*
* This source code is licensed under the CC-by-NC license found in the
* LICENSE file in the root directory of this source tree.
*/
// Copyright 2004-present Facebook. All Rights Reserved.
#include "WorkerThread.h"
#include "../../FaissAssert.h"
namespace faiss { namespace gpu {
WorkerThread::WorkerThread() :
wantStop_(false) {
startThread();
// Make sure that the thread has started before continuing
add([](){}).get();
}
WorkerThread::~WorkerThread() {
stop();
waitForThreadExit();
}
void
WorkerThread::startThread() {
thread_ = std::thread([this](){ threadMain(); });
}
void
WorkerThread::stop() {
std::lock_guard<std::mutex> guard(mutex_);
wantStop_ = true;
monitor_.notify_one();
}
std::future<bool>
WorkerThread::add(std::function<void()> f) {
std::lock_guard<std::mutex> guard(mutex_);
if (wantStop_) {
// The timer thread has been stopped, or we want to stop; we can't
// schedule anything else
std::promise<bool> p;
auto fut = p.get_future();
// did not execute
p.set_value(false);
return fut;
}
auto pr = std::promise<bool>();
auto fut = pr.get_future();
queue_.emplace_back(std::make_pair(std::move(f), std::move(pr)));
// Wake up our thread
monitor_.notify_one();
return fut;
}
void
WorkerThread::threadMain() {
threadLoop();
// Call all pending tasks
FAISS_ASSERT(wantStop_);
for (auto& f : queue_) {
f.first();
f.second.set_value(true);
}
}
void
WorkerThread::threadLoop() {
while (true) {
std::pair<std::function<void()>, std::promise<bool>> data;
{
std::unique_lock<std::mutex> lock(mutex_);
while (!wantStop_ && queue_.empty()) {
monitor_.wait(lock);
}
if (wantStop_) {
return;
}
data = std::move(queue_.front());
queue_.pop_front();
}
data.first();
data.second.set_value(true);
}
}
void
WorkerThread::waitForThreadExit() {
try {
thread_.join();
} catch (...) {
}
}
} } // namespace