/Users/deen/code/yugabyte-db/src/yb/util/test_thread_holder.h
Line | Count | Source (jump to first uncovered line) |
1 | | // Copyright (c) YugaByte, Inc. |
2 | | // |
3 | | // Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except |
4 | | // in compliance with the License. You may obtain a copy of the License at |
5 | | // |
6 | | // http://www.apache.org/licenses/LICENSE-2.0 |
7 | | // |
8 | | // Unless required by applicable law or agreed to in writing, software distributed under the License |
9 | | // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express |
10 | | // or implied. See the License for the specific language governing permissions and limitations |
11 | | // under the License. |
12 | | // |
13 | | |
14 | | #ifndef YB_UTIL_TEST_THREAD_HOLDER_H |
15 | | #define YB_UTIL_TEST_THREAD_HOLDER_H |
16 | | |
17 | | #include <thread> |
18 | | |
19 | | #include "yb/util/monotime.h" |
20 | | #include "yb/util/status.h" |
21 | | #include "yb/util/thread.h" |
22 | | |
23 | | namespace yb { |
24 | | |
25 | | class SetFlagOnExit { |
26 | | public: |
27 | | explicit SetFlagOnExit(std::atomic<bool>* stop_flag) |
28 | 155 | : stop_flag_(stop_flag) {} |
29 | | |
30 | 153 | ~SetFlagOnExit() { |
31 | 153 | stop_flag_->store(true, std::memory_order_release); |
32 | 153 | } |
33 | | |
34 | | private: |
35 | | std::atomic<bool>* stop_flag_; |
36 | | }; |
37 | | |
38 | | // Waits specified duration or when stop switches to true. |
39 | | void WaitStopped(const CoarseDuration& duration, std::atomic<bool>* stop); |
40 | | |
41 | | // Holds vector of threads, and provides convenient utilities. Such as JoinAll, Wait etc. |
42 | | class TestThreadHolder { |
43 | | public: |
44 | | ~TestThreadHolder() { |
45 | | stop_flag_.store(true, std::memory_order_release); |
46 | | JoinAll(); |
47 | | } |
48 | | |
49 | | template <class... Args> |
50 | | void AddThread(Args&&... args) { |
51 | | threads_.emplace_back(std::forward<Args>(args)...); |
52 | | } |
53 | | |
54 | 0 | void AddThread(std::thread thread) { |
55 | 0 | threads_.push_back(std::move(thread)); |
56 | 0 | } |
57 | | |
58 | | template <class Functor> |
59 | | void AddThreadFunctor(const Functor& functor) { |
60 | | AddThread([&stop = stop_flag_, functor] { |
61 | | CDSAttacher attacher; |
62 | | SetFlagOnExit set_stop_on_exit(&stop); |
63 | | functor(); |
64 | | }); |
65 | | } |
66 | | |
67 | 0 | void Wait(const CoarseDuration& duration) { |
68 | 0 | WaitStopped(duration, &stop_flag_); |
69 | 0 | } |
70 | | |
71 | | void JoinAll(); |
72 | | |
73 | | template <class Cond> |
74 | | CHECKED_STATUS WaitCondition(const Cond& cond) { |
75 | | while (!cond()) { |
76 | | if (stop_flag_.load(std::memory_order_acquire)) { |
77 | | return STATUS(Aborted, "Wait aborted"); |
78 | | } |
79 | | std::this_thread::sleep_for(std::chrono::milliseconds(100)); |
80 | | } |
81 | | |
82 | | return Status::OK(); |
83 | | } |
84 | | |
85 | | void WaitAndStop(const CoarseDuration& duration) { |
86 | | yb::WaitStopped(duration, &stop_flag_); |
87 | | Stop(); |
88 | | } |
89 | | |
90 | | void Stop() { |
91 | | stop_flag_.store(true, std::memory_order_release); |
92 | | JoinAll(); |
93 | | } |
94 | | |
95 | | std::atomic<bool>& stop_flag() { |
96 | | return stop_flag_; |
97 | | } |
98 | | |
99 | | private: |
100 | | std::atomic<bool> stop_flag_{false}; |
101 | | std::vector<std::thread> threads_; |
102 | | }; |
103 | | |
104 | | } // namespace yb |
105 | | |
106 | | #endif // YB_UTIL_TEST_THREAD_HOLDER_H |