folly copyright 2015 -> copyright 2016
[folly.git] / folly / experimental / RCURefCount.h
1 /*
2  * Copyright 2016 Facebook, Inc.
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 #pragma once
17
18 #include <folly/ThreadLocal.h>
19 #include <folly/experimental/RCUUtils.h>
20
21 namespace folly {
22
23 class RCURefCount {
24  public:
25   using Int = int64_t;
26
27   RCURefCount() :
28       localCount_([&]() {
29           return new LocalRefCount(globalCount_);
30         }) {
31   }
32
33   ~RCURefCount() noexcept {
34     assert(state_ == State::GLOBAL);
35     assert(globalCount_.load() == 0);
36   }
37
38   // This can't increment from 0.
39   Int operator++() noexcept {
40     auto& localCount = *localCount_;
41
42     std::lock_guard<RCUReadLock> lg(RCUReadLock::instance());
43     auto state = state_.load();
44
45     if (LIKELY(state == State::LOCAL)) {
46       ++localCount;
47
48       return 42;
49     } else if (state == State::GLOBAL_TRANSITION) {
50       ++globalCount_;
51
52       return 42;
53     } else {
54       auto globalCount = globalCount_.load();
55
56       do {
57         if (!globalCount) {
58           return 0;
59         }
60       } while (!globalCount_.compare_exchange_weak(globalCount,
61                                                    globalCount + 1));
62
63       return globalCount + 1;
64     }
65   }
66
67   Int operator--() noexcept {
68     auto& localCount = *localCount_;
69
70     std::lock_guard<RCUReadLock> lg(RCUReadLock::instance());
71     auto state = state_.load();
72
73     if (LIKELY(state == State::LOCAL)) {
74       --localCount;
75
76       return 42;
77     } else {
78       auto value = --globalCount_;
79
80       if (state == State::GLOBAL) {
81         assert(value >= 0);
82         return value;
83       } else {
84         return 42;
85       }
86     }
87   }
88
89   Int operator*() const {
90     std::lock_guard<RCUReadLock> lg(RCUReadLock::instance());
91
92     if (state_ == State::GLOBAL) {
93       return globalCount_;
94     }
95
96     return 42;
97   }
98
99   void useGlobal() noexcept {
100     state_ = State::GLOBAL_TRANSITION;
101
102     synchronize_rcu();
103     // At this point everyone is using the global count
104
105     auto accessor = localCount_.accessAllThreads();
106     for (auto& count : accessor) {
107       count.collect();
108     }
109
110     state_ = State::GLOBAL;
111
112     synchronize_rcu();
113     // After this ++ or -- can return 0.
114   }
115
116  private:
117   using AtomicInt = std::atomic<Int>;
118
119   enum class State {
120     LOCAL,
121     GLOBAL_TRANSITION,
122     GLOBAL
123   };
124
125   class LocalRefCount {
126    public:
127     explicit LocalRefCount(AtomicInt& globalCount) :
128         count_(0),
129         globalCount_(globalCount) {
130       RCURegisterThread();
131     }
132
133     ~LocalRefCount() {
134       collect();
135     }
136
137     void collect() {
138       globalCount_ += count_;
139       count_ = 0;
140     }
141
142     void operator++() {
143       ++count_;
144     }
145
146     void operator--() {
147       --count_;
148     }
149
150    private:
151     Int count_;
152     AtomicInt& globalCount_;
153   };
154
155   std::atomic<State> state_{State::LOCAL};
156   folly::ThreadLocal<LocalRefCount, RCURefCount> localCount_;
157   std::atomic<int64_t> globalCount_{1};
158 };
159
160 }