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_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 "nsync_cv.h" 26 #include "nsync_mu.h" 27 #include "tensorflow/core/platform/thread_annotations.h" 28 namespace tensorflow { 29 30 #undef mutex_lock 31 32 enum LinkerInitialized { LINKER_INITIALIZED }; 33 34 class condition_variable; 35 36 // Mimic std::mutex + C++17's shared_mutex, adding a LinkerInitialized 37 // constructor interface. This type is as fast as mutex, but is also a shared 38 // lock. 39 class LOCKABLE mutex { 40 public: 41 mutex() { nsync::nsync_mu_init(&mu_); } 42 // The default implementation of nsync_mutex is safe to use after the linker 43 // initializations 44 explicit mutex(LinkerInitialized x) {} 45 46 void lock() EXCLUSIVE_LOCK_FUNCTION() { nsync::nsync_mu_lock(&mu_); } 47 bool try_lock() EXCLUSIVE_TRYLOCK_FUNCTION(true) { 48 return nsync::nsync_mu_trylock(&mu_) != 0; 49 }; 50 void unlock() UNLOCK_FUNCTION() { nsync::nsync_mu_unlock(&mu_); } 51 52 void lock_shared() SHARED_LOCK_FUNCTION() { nsync::nsync_mu_rlock(&mu_); } 53 bool try_lock_shared() SHARED_TRYLOCK_FUNCTION(true) { 54 return nsync::nsync_mu_rtrylock(&mu_) != 0; 55 }; 56 void unlock_shared() UNLOCK_FUNCTION() { nsync::nsync_mu_runlock(&mu_); } 57 58 private: 59 friend class condition_variable; 60 nsync::nsync_mu mu_; 61 }; 62 63 // Mimic a subset of the std::unique_lock<tensorflow::mutex> functionality. 64 class SCOPED_LOCKABLE mutex_lock { 65 public: 66 typedef ::tensorflow::mutex mutex_type; 67 68 explicit mutex_lock(mutex_type& mu) EXCLUSIVE_LOCK_FUNCTION(mu) : mu_(&mu) { 69 mu_->lock(); 70 } 71 72 mutex_lock(mutex_type& mu, std::try_to_lock_t) EXCLUSIVE_LOCK_FUNCTION(mu) 73 : mu_(&mu) { 74 if (!mu.try_lock()) { 75 mu_ = nullptr; 76 } 77 } 78 79 // Manually nulls out the source to prevent double-free. 80 // (std::move does not null the source pointer by default.) 81 explicit mutex_lock(mutex_lock&& ml) noexcept : mu_(ml.mu_) { 82 ml.mu_ = nullptr; 83 } 84 ~mutex_lock() UNLOCK_FUNCTION() { 85 if (mu_ != nullptr) { 86 mu_->unlock(); 87 } 88 } 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 106 explicit tf_shared_lock(mutex_type& mu) SHARED_LOCK_FUNCTION(mu) : mu_(&mu) { 107 mu_->lock_shared(); 108 } 109 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.) 119 explicit tf_shared_lock(tf_shared_lock&& ml) noexcept : mu_(ml.mu_) { 120 ml.mu_ = nullptr; 121 } 122 ~tf_shared_lock() UNLOCK_FUNCTION() { 123 if (mu_ != nullptr) { 124 mu_->unlock_shared(); 125 } 126 } 127 mutex_type* mutex() { return mu_; } 128 129 operator bool() const { return mu_ != nullptr; } 130 131 private: 132 mutex_type* mu_; 133 }; 134 135 // Catch bug where variable name is omitted, e.g. tf_shared_lock (mu); 136 #define tf_shared_lock(x) \ 137 static_assert(0, "tf_shared_lock_decl_missing_var_name"); 138 139 // Mimic std::condition_variable. 140 class condition_variable { 141 public: 142 condition_variable() { nsync::nsync_cv_init(&cv_); } 143 144 void wait(mutex_lock& lock) { 145 nsync::nsync_cv_wait(&cv_, &lock.mutex()->mu_); 146 } 147 template <class Rep, class Period> 148 std::cv_status wait_for(mutex_lock& lock, 149 std::chrono::duration<Rep, Period> dur) { 150 int r = nsync::nsync_cv_wait_with_deadline( 151 &cv_, &lock.mutex()->mu_, std::chrono::system_clock::now() + dur, 152 nullptr); 153 return r ? std::cv_status::timeout : std::cv_status::no_timeout; 154 } 155 void notify_one() { nsync::nsync_cv_signal(&cv_); } 156 void notify_all() { nsync::nsync_cv_broadcast(&cv_); } 157 158 private: 159 friend ConditionResult WaitForMilliseconds(mutex_lock* mu, 160 condition_variable* cv, int64 ms); 161 nsync::nsync_cv cv_; 162 }; 163 164 inline ConditionResult WaitForMilliseconds(mutex_lock* mu, 165 condition_variable* cv, int64 ms) { 166 std::cv_status s = cv->wait_for(*mu, std::chrono::milliseconds(ms)); 167 return (s == std::cv_status::timeout) ? kCond_Timeout : kCond_MaybeNotified; 168 } 169 170 } // namespace tensorflow 171 172 #endif // TENSORFLOW_PLATFORM_DEFAULT_MUTEX_H_ 173
