Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
WorkerThread.cpp
1 /**
2 * Copyright (c) 2015-present, Facebook, Inc.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD+Patents license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9
10 #include "WorkerThread.h"
11 #include "../../FaissAssert.h"
12
13 namespace faiss { namespace gpu {
14
15 WorkerThread::WorkerThread() :
16 wantStop_(false) {
17 startThread();
18
19 // Make sure that the thread has started before continuing
20 add([](){}).get();
21 }
22
23 WorkerThread::~WorkerThread() {
24 stop();
25 waitForThreadExit();
26 }
27
28 void
29 WorkerThread::startThread() {
30 thread_ = std::thread([this](){ threadMain(); });
31 }
32
33 void
34 WorkerThread::stop() {
35 std::lock_guard<std::mutex> guard(mutex_);
36
37 wantStop_ = true;
38 monitor_.notify_one();
39 }
40
41 std::future<bool>
42 WorkerThread::add(std::function<void()> f) {
43 std::lock_guard<std::mutex> guard(mutex_);
44
45 if (wantStop_) {
46 // The timer thread has been stopped, or we want to stop; we can't
47 // schedule anything else
48 std::promise<bool> p;
49 auto fut = p.get_future();
50
51 // did not execute
52 p.set_value(false);
53 return fut;
54 }
55
56 auto pr = std::promise<bool>();
57 auto fut = pr.get_future();
58
59 queue_.emplace_back(std::make_pair(std::move(f), std::move(pr)));
60
61 // Wake up our thread
62 monitor_.notify_one();
63 return fut;
64 }
65
66 void
67 WorkerThread::threadMain() {
68 threadLoop();
69
70 // Call all pending tasks
71 FAISS_ASSERT(wantStop_);
72
73 for (auto& f : queue_) {
74 f.first();
75 f.second.set_value(true);
76 }
77 }
78
79 void
80 WorkerThread::threadLoop() {
81 while (true) {
82 std::pair<std::function<void()>, std::promise<bool>> data;
83
84 {
85 std::unique_lock<std::mutex> lock(mutex_);
86
87 while (!wantStop_ && queue_.empty()) {
88 monitor_.wait(lock);
89 }
90
91 if (wantStop_) {
92 return;
93 }
94
95 data = std::move(queue_.front());
96 queue_.pop_front();
97 }
98
99 data.first();
100 data.second.set_value(true);
101 }
102 }
103
104 void
105 WorkerThread::waitForThreadExit() {
106 try {
107 thread_.join();
108 } catch (...) {
109 }
110 }
111
112 } } // namespace
faiss::gpu::WorkerThread::stop
void stop()
Request that the worker thread stop itself.
Definition: WorkerThread.cpp:34
faiss::gpu::WorkerThread::add
std::future< bool > add(std::function< void()> f)
Definition: WorkerThread.cpp:42
faiss::gpu::WorkerThread::waitForThreadExit
void waitForThreadExit()
Definition: WorkerThread.cpp:105
faiss::gpu::WorkerThread::~WorkerThread
~WorkerThread()
Definition: WorkerThread.cpp:23