2 * Copyright 2015 Facebook, Inc.
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
8 * http://www.apache.org/licenses/LICENSE-2.0
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.
18 #include <folly/Baton.h>
19 #include <folly/ThreadLocal.h>
29 return new LocalRefCount(*this);
31 collectGuard_(&collectBaton_, [](void* p) {
32 auto baton = reinterpret_cast<folly::Baton<>*>(p);
37 ~TLRefCount() noexcept {
38 assert(globalCount_.load() == 0);
39 assert(state_.load() == State::GLOBAL);
42 // This can't increment from 0.
43 Int operator++() noexcept {
44 auto& localCount = *localCount_;
50 if (state_.load() == State::GLOBAL_TRANSITION) {
51 std::lock_guard<std::mutex> lg(globalMutex_);
54 assert(state_.load() == State::GLOBAL);
56 auto value = globalCount_.load();
61 } while (!globalCount_.compare_exchange_weak(value, value+1));
66 Int operator--() noexcept {
67 auto& localCount = *localCount_;
73 if (state_.load() == State::GLOBAL_TRANSITION) {
74 std::lock_guard<std::mutex> lg(globalMutex_);
77 assert(state_.load() == State::GLOBAL);
79 return --globalCount_;
82 Int operator*() const {
83 if (state_ != State::GLOBAL) {
86 return globalCount_.load();
89 void useGlobal() noexcept {
90 std::lock_guard<std::mutex> lg(globalMutex_);
92 state_ = State::GLOBAL_TRANSITION;
94 auto accessor = localCount_.accessAllThreads();
95 for (auto& count : accessor) {
99 collectGuard_.reset();
100 collectBaton_.wait();
102 state_ = State::GLOBAL;
106 using AtomicInt = std::atomic<Int>;
114 class LocalRefCount {
116 explicit LocalRefCount(TLRefCount& refCount) :
117 refCount_(refCount) {
118 std::lock_guard<std::mutex> lg(refCount.globalMutex_);
120 collectGuard_ = refCount.collectGuard_;
128 std::lock_guard<std::mutex> lg(collectMutex_);
130 if (!collectGuard_) {
134 collectCount_ = count_;
135 refCount_.globalCount_ += collectCount_;
136 collectGuard_.reset();
148 bool update(Int delta) {
149 if (UNLIKELY(refCount_.state_.load() != State::LOCAL)) {
153 auto count = count_ += delta;
155 if (UNLIKELY(refCount_.state_.load() != State::LOCAL)) {
156 std::lock_guard<std::mutex> lg(collectMutex_);
161 if (collectCount_ != count) {
170 TLRefCount& refCount_;
172 std::mutex collectMutex_;
173 Int collectCount_{0};
174 std::shared_ptr<void> collectGuard_;
177 std::atomic<State> state_{State::LOCAL};
178 folly::ThreadLocal<LocalRefCount, TLRefCount> localCount_;
179 std::atomic<int64_t> globalCount_{1};
180 std::mutex globalMutex_;
181 folly::Baton<> collectBaton_;
182 std::shared_ptr<void> collectGuard_;