1 /** 2 * Copyright 2019 Huawei Technologies Co., Ltd 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 #ifndef MINDSPORE_CCSRC_MINDDATA_DATASET_UTIL_LOCK_H_ 17 #define MINDSPORE_CCSRC_MINDDATA_DATASET_UTIL_LOCK_H_ 18 19 #include <atomic> 20 #include <condition_variable> 21 #include <mutex> 22 23 namespace mindspore { 24 namespace dataset { 25 class SpinLock { 26 public: 27 void Lock(); 28 29 bool TryLock(); 30 31 void Unlock() noexcept; 32 SpinLock()33 SpinLock() : val_(kUnlocked) {} 34 35 SpinLock(const SpinLock &) = delete; 36 37 SpinLock(SpinLock &&) = delete; 38 39 ~SpinLock() = default; 40 41 SpinLock &operator=(const SpinLock &) = delete; 42 43 SpinLock &operator=(SpinLock &&) = delete; 44 45 private: 46 static constexpr int kUnlocked = 0; 47 static constexpr int kLocked = 1; 48 std::atomic<int> val_; 49 }; 50 51 // C++11 has no shared mutex. The following class is an alternative. It favors writer and is suitable for the case 52 // where writer is rare. 53 class RWLock { 54 public: RWLock()55 RWLock() : status_(0), waiting_readers_(0), waiting_writers_(0) {} 56 57 RWLock(const RWLock &) = delete; 58 59 RWLock(RWLock &&) = delete; 60 61 ~RWLock() = default; 62 63 RWLock &operator=(const RWLock &) = delete; 64 65 RWLock &operator=(RWLock &&) = delete; 66 67 void LockShared(); 68 LockExclusive()69 void LockExclusive() { 70 std::unique_lock<std::mutex> lck(mtx_); 71 waiting_writers_ += 1; 72 write_cv_.wait(lck, [this]() { return status_ == 0; }); 73 waiting_writers_ -= 1; 74 status_ = -1; 75 } 76 77 void Unlock() noexcept; 78 79 // Upgrade a shared lock to exclusive lock 80 void Upgrade(); 81 82 // Downgrade an exclusive lock to shared lock 83 void Downgrade(); 84 85 private: 86 // -1 : one writer 87 // 0 : no reader and no writer 88 // n > 0 : n reader 89 int32_t status_; 90 int32_t waiting_readers_; 91 int32_t waiting_writers_; 92 std::mutex mtx_; 93 std::condition_variable read_cv_; 94 std::condition_variable write_cv_; 95 }; 96 97 // A Wrapper for RWLock. The destructor will release the lock if we own it. 98 class SharedLock { 99 public: 100 explicit SharedLock(RWLock *rw); 101 102 ~SharedLock(); 103 104 SharedLock(const SharedLock &) = delete; 105 106 SharedLock(SharedLock &&) = delete; 107 108 SharedLock &operator=(const SharedLock &) = delete; 109 110 SharedLock &operator=(SharedLock &&) = delete; 111 112 void Unlock(); 113 114 void Lock(); 115 116 void Upgrade(); 117 118 void Downgrade(); 119 120 private: 121 RWLock *rw_; 122 bool ownlock_; 123 }; 124 125 class UniqueLock { 126 public: 127 explicit UniqueLock(RWLock *rw); 128 129 ~UniqueLock(); 130 131 UniqueLock(const UniqueLock &) = delete; 132 133 UniqueLock(UniqueLock &&) = delete; 134 135 UniqueLock &operator=(const UniqueLock &) = delete; 136 137 UniqueLock &operator=(UniqueLock &&) = delete; 138 139 void Unlock(); 140 141 void Lock(); 142 143 private: 144 RWLock *rw_; 145 bool ownlock_; 146 }; 147 148 class LockGuard { 149 public: 150 explicit LockGuard(SpinLock *lock); 151 152 ~LockGuard(); 153 154 LockGuard(const LockGuard &) = delete; 155 156 LockGuard(LockGuard &&) = delete; 157 158 LockGuard &operator=(const LockGuard &) = delete; 159 160 LockGuard &operator=(LockGuard &&) = delete; 161 162 void Unlock(); 163 164 void Lock(); 165 166 private: 167 SpinLock *lck_; 168 bool own_lock_; 169 }; 170 } // namespace dataset 171 } // namespace mindspore 172 173 #endif // MINDSPORE_CCSRC_MINDDATA_DATASET_UTIL_LOCK_H_ 174