1 /* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15
16 #ifndef TENSORFLOW_CORE_PLATFORM_DEFAULT_MUTEX_H_
17 #define TENSORFLOW_CORE_PLATFORM_DEFAULT_MUTEX_H_
18
19 // IWYU pragma: private, include "third_party/tensorflow/core/platform/mutex.h"
20 // IWYU pragma: friend third_party/tensorflow/core/platform/mutex.h
21
22 #include <chrono>
23 #include <condition_variable>
24 #include <mutex>
25 #include "tensorflow/core/platform/thread_annotations.h"
26
27 namespace tensorflow {
28
29 #undef mutex_lock
30
31 enum LinkerInitialized { LINKER_INITIALIZED };
32
33 class condition_variable;
34
35 // Mimic std::mutex + C++17's shared_mutex, adding a LinkerInitialized
36 // constructor interface. This type is as fast as mutex, but is also a shared
37 // lock.
38 class LOCKABLE mutex {
39 public:
40 mutex();
41 // The default implementation of the underlying mutex is safe to use after
42 // the linker initialization to zero.
mutex(LinkerInitialized x)43 explicit mutex(LinkerInitialized x) {}
44
45 void lock() EXCLUSIVE_LOCK_FUNCTION();
46 bool try_lock() EXCLUSIVE_TRYLOCK_FUNCTION(true);
47 void unlock() UNLOCK_FUNCTION();
48
49 void lock_shared() SHARED_LOCK_FUNCTION();
50 bool try_lock_shared() SHARED_TRYLOCK_FUNCTION(true);
51 void unlock_shared() UNLOCK_FUNCTION();
52
53 struct external_mu_space {
54 void* space[2];
55 };
56
57 private:
58 friend class condition_variable;
59 external_mu_space mu_;
60 };
61
62 // Mimic a subset of the std::unique_lock<tensorflow::mutex> functionality.
63 class SCOPED_LOCKABLE mutex_lock {
64 public:
65 typedef ::tensorflow::mutex mutex_type;
66
mutex_lock(mutex_type & mu)67 explicit mutex_lock(mutex_type& mu) EXCLUSIVE_LOCK_FUNCTION(mu) : mu_(&mu) {
68 mu_->lock();
69 }
70
mutex_lock(mutex_type & mu,std::try_to_lock_t)71 mutex_lock(mutex_type& mu, std::try_to_lock_t) EXCLUSIVE_LOCK_FUNCTION(mu)
72 : mu_(&mu) {
73 if (!mu.try_lock()) {
74 mu_ = nullptr;
75 }
76 }
77
78 // Manually nulls out the source to prevent double-free.
79 // (std::move does not null the source pointer by default.)
mutex_lock(mutex_lock && ml)80 mutex_lock(mutex_lock&& ml) noexcept EXCLUSIVE_LOCK_FUNCTION(ml.mu_)
81 : mu_(ml.mu_) {
82 ml.mu_ = nullptr;
83 }
UNLOCK_FUNCTION()84 ~mutex_lock() UNLOCK_FUNCTION() {
85 if (mu_ != nullptr) {
86 mu_->unlock();
87 }
88 }
mutex()89 mutex_type* mutex() { return mu_; }
90
91 operator bool() const { return mu_ != nullptr; }
92
93 private:
94 mutex_type* mu_;
95 };
96
97 // Catch bug where variable name is omitted, e.g. mutex_lock (mu);
98 #define mutex_lock(x) static_assert(0, "mutex_lock_decl_missing_var_name");
99
100 // Mimic a subset of the std::shared_lock<tensorflow::mutex> functionality.
101 // Name chosen to minimise conflicts with the tf_shared_lock macro, below.
102 class SCOPED_LOCKABLE tf_shared_lock {
103 public:
104 typedef ::tensorflow::mutex mutex_type;
105
tf_shared_lock(mutex_type & mu)106 explicit tf_shared_lock(mutex_type& mu) SHARED_LOCK_FUNCTION(mu) : mu_(&mu) {
107 mu_->lock_shared();
108 }
109
tf_shared_lock(mutex_type & mu,std::try_to_lock_t)110 tf_shared_lock(mutex_type& mu, std::try_to_lock_t) SHARED_LOCK_FUNCTION(mu)
111 : mu_(&mu) {
112 if (!mu.try_lock_shared()) {
113 mu_ = nullptr;
114 }
115 }
116
117 // Manually nulls out the source to prevent double-free.
118 // (std::move does not null the source pointer by default.)
tf_shared_lock(tf_shared_lock && ml)119 tf_shared_lock(tf_shared_lock&& ml) noexcept SHARED_LOCK_FUNCTION(ml.mu_)
120 : mu_(ml.mu_) {
121 ml.mu_ = nullptr;
122 }
UNLOCK_FUNCTION()123 ~tf_shared_lock() UNLOCK_FUNCTION() {
124 if (mu_ != nullptr) {
125 mu_->unlock_shared();
126 }
127 }
mutex()128 mutex_type* mutex() { return mu_; }
129
130 operator bool() const { return mu_ != nullptr; }
131
132 private:
133 mutex_type* mu_;
134 };
135
136 // Catch bug where variable name is omitted, e.g. tf_shared_lock (mu);
137 #define tf_shared_lock(x) \
138 static_assert(0, "tf_shared_lock_decl_missing_var_name");
139
140 // Mimic std::condition_variable.
141 class condition_variable {
142 public:
143 condition_variable();
144
145 void wait(mutex_lock& lock);
146 template <class Rep, class Period>
wait_for(mutex_lock & lock,std::chrono::duration<Rep,Period> dur)147 std::cv_status wait_for(mutex_lock& lock,
148 std::chrono::duration<Rep, Period> dur) {
149 return wait_until_system_clock(lock,
150 std::chrono::system_clock::now() + dur);
151 }
152 void notify_one();
153 void notify_all();
154
155 struct external_cv_space {
156 void* space[2];
157 };
158
159 private:
160 friend ConditionResult WaitForMilliseconds(mutex_lock* mu,
161 condition_variable* cv, int64 ms);
162 std::cv_status wait_until_system_clock(
163 mutex_lock& lock,
164 const std::chrono::system_clock::time_point timeout_time);
165 external_cv_space cv_;
166 };
167
WaitForMilliseconds(mutex_lock * mu,condition_variable * cv,int64 ms)168 inline ConditionResult WaitForMilliseconds(mutex_lock* mu,
169 condition_variable* cv, int64 ms) {
170 std::cv_status s = cv->wait_for(*mu, std::chrono::milliseconds(ms));
171 return (s == std::cv_status::timeout) ? kCond_Timeout : kCond_MaybeNotified;
172 }
173
174 } // namespace tensorflow
175
176 #endif // TENSORFLOW_CORE_PLATFORM_DEFAULT_MUTEX_H_
177