1/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations under the License.
14==============================================================================*/
15
16#ifndef TENSORFLOW_PLATFORM_DEFAULT_MUTEX_H_
17#define TENSORFLOW_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
27namespace tensorflow {
28
29#undef mutex_lock
30
31enum LinkerInitialized { LINKER_INITIALIZED };
32
33class 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.
38class 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.
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.
63class SCOPED_LOCKABLE mutex_lock {
64 public:
65 typedef ::tensorflow::mutex mutex_type;
66
67 explicit mutex_lock(mutex_type& mu) EXCLUSIVE_LOCK_FUNCTION(mu) : mu_(&mu) {
68 mu_->lock();
69 }
70
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.)
80 explicit mutex_lock(mutex_lock&& ml) noexcept : mu_(ml.mu_) {
81 ml.mu_ = nullptr;
82 }
83 ~mutex_lock() UNLOCK_FUNCTION() {
84 if (mu_ != nullptr) {
85 mu_->unlock();
86 }
87 }
88 mutex_type* mutex() { return mu_; }
89
90 operator bool() const { return mu_ != nullptr; }
91
92 private:
93 mutex_type* mu_;
94};
95
96// Catch bug where variable name is omitted, e.g. mutex_lock (mu);
97#define mutex_lock(x) static_assert(0, "mutex_lock_decl_missing_var_name");
98
99// Mimic a subset of the std::shared_lock<tensorflow::mutex> functionality.
100// Name chosen to minimise conflicts with the tf_shared_lock macro, below.
101class SCOPED_LOCKABLE tf_shared_lock {
102 public:
103 typedef ::tensorflow::mutex mutex_type;
104
105 explicit tf_shared_lock(mutex_type& mu) SHARED_LOCK_FUNCTION(mu) : mu_(&mu) {
106 mu_->lock_shared();
107 }
108
109 tf_shared_lock(mutex_type& mu, std::try_to_lock_t) SHARED_LOCK_FUNCTION(mu)
110 : mu_(&mu) {
111 if (!mu.try_lock_shared()) {
112 mu_ = nullptr;
113 }
114 }
115
116 // Manually nulls out the source to prevent double-free.
117 // (std::move does not null the source pointer by default.)
118 explicit tf_shared_lock(tf_shared_lock&& ml) noexcept : mu_(ml.mu_) {
119 ml.mu_ = nullptr;
120 }
121 ~tf_shared_lock() UNLOCK_FUNCTION() {
122 if (mu_ != nullptr) {
123 mu_->unlock_shared();
124 }
125 }
126 mutex_type* mutex() { return mu_; }
127
128 operator bool() const { return mu_ != nullptr; }
129
130 private:
131 mutex_type* mu_;
132};
133
134// Catch bug where variable name is omitted, e.g. tf_shared_lock (mu);
135#define tf_shared_lock(x) \
136 static_assert(0, "tf_shared_lock_decl_missing_var_name");
137
138// Mimic std::condition_variable.
139class condition_variable {
140 public:
141 condition_variable();
142
143 void wait(mutex_lock& lock);
144 template <class Rep, class Period>
145 std::cv_status wait_for(mutex_lock& lock,
146 std::chrono::duration<Rep, Period> dur) {
147 return wait_until_system_clock(lock,
148 std::chrono::system_clock::now() + dur);
149 }
150 void notify_one();
151 void notify_all();
152
153 struct external_cv_space {
154 void* space[2];
155 };
156
157 private:
158 friend ConditionResult WaitForMilliseconds(mutex_lock* mu,
159 condition_variable* cv, int64 ms);
160 std::cv_status wait_until_system_clock(
161 mutex_lock& lock,
162 const std::chrono::system_clock::time_point timeout_time);
163 external_cv_space cv_;
164};
165
166inline ConditionResult WaitForMilliseconds(mutex_lock* mu,
167 condition_variable* cv, int64 ms) {
168 std::cv_status s = cv->wait_for(*mu, std::chrono::milliseconds(ms));
169 return (s == std::cv_status::timeout) ? kCond_Timeout : kCond_MaybeNotified;
170}
171
172} // namespace tensorflow
173
174#endif // TENSORFLOW_PLATFORM_DEFAULT_MUTEX_H_
175