Add prepareSkipTo() method to EliasFanoReader
[folly.git] / folly / ConcurrentSkipList-inl.h
1 /*
2  * Copyright 2017 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
17 // @author: Xin Liu <xliux@fb.com>
18
19 #pragma once
20
21 #include <algorithm>
22 #include <atomic>
23 #include <climits>
24 #include <cmath>
25 #include <memory>
26 #include <mutex>
27 #include <type_traits>
28 #include <vector>
29
30 #include <boost/noncopyable.hpp>
31 #include <boost/random.hpp>
32 #include <boost/type_traits.hpp>
33 #include <glog/logging.h>
34
35 #include <folly/Memory.h>
36 #include <folly/MicroSpinLock.h>
37 #include <folly/ThreadLocal.h>
38
39 namespace folly { namespace detail {
40
41 template <typename ValT, typename NodeT> class csl_iterator;
42
43 template <typename T>
44 class SkipListNode : private boost::noncopyable {
45   enum : uint16_t {
46     IS_HEAD_NODE = 1,
47     MARKED_FOR_REMOVAL = (1 << 1),
48     FULLY_LINKED = (1 << 2),
49   };
50
51  public:
52   typedef T value_type;
53
54   template <
55       typename NodeAlloc,
56       typename U,
57       typename =
58           typename std::enable_if<std::is_convertible<U, T>::value>::type>
59   static SkipListNode* create(
60       NodeAlloc& alloc, int height, U&& data, bool isHead = false) {
61     DCHECK(height >= 1 && height < 64) << height;
62
63     size_t size = sizeof(SkipListNode) +
64       height * sizeof(std::atomic<SkipListNode*>);
65     auto* node = static_cast<SkipListNode*>(alloc.allocate(size));
66     // do placement new
67     new (node) SkipListNode(uint8_t(height), std::forward<U>(data), isHead);
68     return node;
69   }
70
71   template <typename NodeAlloc>
72   static void destroy(NodeAlloc& alloc, SkipListNode* node) {
73     node->~SkipListNode();
74     alloc.deallocate(node);
75   }
76
77   template <typename NodeAlloc>
78   struct DestroyIsNoOp : std::integral_constant<bool,
79     IsArenaAllocator<NodeAlloc>::value &&
80     boost::has_trivial_destructor<SkipListNode>::value> { };
81
82   // copy the head node to a new head node assuming lock acquired
83   SkipListNode* copyHead(SkipListNode* node) {
84     DCHECK(node != nullptr && height_ > node->height_);
85     setFlags(node->getFlags());
86     for (uint8_t i = 0; i < node->height_; ++i) {
87       setSkip(i, node->skip(i));
88     }
89     return this;
90   }
91
92   inline SkipListNode* skip(int layer) const {
93     DCHECK_LT(layer, height_);
94     return skip_[layer].load(std::memory_order_consume);
95   }
96
97   // next valid node as in the linked list
98   SkipListNode* next() {
99     SkipListNode* node;
100     for (node = skip(0);
101         (node != nullptr && node->markedForRemoval());
102         node = node->skip(0)) {}
103     return node;
104   }
105
106   void setSkip(uint8_t h, SkipListNode* next) {
107     DCHECK_LT(h, height_);
108     skip_[h].store(next, std::memory_order_release);
109   }
110
111   value_type& data() { return data_; }
112   const value_type& data() const { return data_; }
113   int maxLayer() const { return height_ - 1; }
114   int height() const { return height_; }
115
116   std::unique_lock<MicroSpinLock> acquireGuard() {
117     return std::unique_lock<MicroSpinLock>(spinLock_);
118   }
119
120   bool fullyLinked() const      { return getFlags() & FULLY_LINKED; }
121   bool markedForRemoval() const { return getFlags() & MARKED_FOR_REMOVAL; }
122   bool isHeadNode() const       { return getFlags() & IS_HEAD_NODE; }
123
124   void setIsHeadNode() {
125     setFlags(uint16_t(getFlags() | IS_HEAD_NODE));
126   }
127   void setFullyLinked() {
128     setFlags(uint16_t(getFlags() | FULLY_LINKED));
129   }
130   void setMarkedForRemoval() {
131     setFlags(uint16_t(getFlags() | MARKED_FOR_REMOVAL));
132   }
133
134  private:
135   // Note! this can only be called from create() as a placement new.
136   template <typename U>
137   SkipListNode(uint8_t height, U&& data, bool isHead) :
138       height_(height), data_(std::forward<U>(data)) {
139     spinLock_.init();
140     setFlags(0);
141     if (isHead) setIsHeadNode();
142     // need to explicitly init the dynamic atomic pointer array
143     for (uint8_t i = 0; i < height_; ++i) {
144       new (&skip_[i]) std::atomic<SkipListNode*>(nullptr);
145     }
146   }
147
148   ~SkipListNode() {
149     for (uint8_t i = 0; i < height_; ++i) {
150       skip_[i].~atomic();
151     }
152   }
153
154   uint16_t getFlags() const {
155     return flags_.load(std::memory_order_consume);
156   }
157   void setFlags(uint16_t flags) {
158     flags_.store(flags, std::memory_order_release);
159   }
160
161   // TODO(xliu): on x86_64, it's possible to squeeze these into
162   // skip_[0] to maybe save 8 bytes depending on the data alignments.
163   // NOTE: currently this is x86_64 only anyway, due to the
164   // MicroSpinLock.
165   std::atomic<uint16_t> flags_;
166   const uint8_t height_;
167   MicroSpinLock spinLock_;
168
169   value_type data_;
170
171   std::atomic<SkipListNode*> skip_[0];
172 };
173
174 class SkipListRandomHeight {
175   enum { kMaxHeight = 64 };
176  public:
177   // make it a singleton.
178   static SkipListRandomHeight *instance() {
179     static SkipListRandomHeight instance_;
180     return &instance_;
181   }
182
183   int getHeight(int maxHeight) const {
184     DCHECK_LE(maxHeight, kMaxHeight) << "max height too big!";
185     double p = randomProb();
186     for (int i = 0; i < maxHeight; ++i) {
187       if (p < lookupTable_[i]) {
188         return i + 1;
189       }
190     }
191     return maxHeight;
192   }
193
194   size_t getSizeLimit(int height) const {
195     DCHECK_LT(height, kMaxHeight);
196     return sizeLimitTable_[height];
197   }
198
199  private:
200   SkipListRandomHeight() { initLookupTable(); }
201
202   void initLookupTable() {
203     // set skip prob = 1/E
204     static const double kProbInv = exp(1);
205     static const double kProb = 1.0 / kProbInv;
206     static const size_t kMaxSizeLimit = std::numeric_limits<size_t>::max();
207
208     double sizeLimit = 1;
209     double p = lookupTable_[0] = (1 - kProb);
210     sizeLimitTable_[0] = 1;
211     for (int i = 1; i < kMaxHeight - 1; ++i) {
212       p *= kProb;
213       sizeLimit *= kProbInv;
214       lookupTable_[i] = lookupTable_[i - 1] + p;
215       sizeLimitTable_[i] = sizeLimit > kMaxSizeLimit ?
216         kMaxSizeLimit :
217         static_cast<size_t>(sizeLimit);
218     }
219     lookupTable_[kMaxHeight - 1] = 1;
220     sizeLimitTable_[kMaxHeight - 1] = kMaxSizeLimit;
221   }
222
223   static double randomProb() {
224     static ThreadLocal<boost::lagged_fibonacci2281> rng_;
225     return (*rng_)();
226   }
227
228   double lookupTable_[kMaxHeight];
229   size_t sizeLimitTable_[kMaxHeight];
230 };
231
232 template <typename NodeType, typename NodeAlloc, typename = void>
233 class NodeRecycler;
234
235 template <typename NodeType, typename NodeAlloc>
236 class NodeRecycler<NodeType, NodeAlloc, typename std::enable_if<
237   !NodeType::template DestroyIsNoOp<NodeAlloc>::value>::type> {
238  public:
239   explicit NodeRecycler(const NodeAlloc& alloc)
240     : refs_(0), dirty_(false), alloc_(alloc) { lock_.init(); }
241
242   explicit NodeRecycler() : refs_(0), dirty_(false) { lock_.init(); }
243
244   ~NodeRecycler() {
245     CHECK_EQ(refs(), 0);
246     if (nodes_) {
247       for (auto& node : *nodes_) {
248         NodeType::destroy(alloc_, node);
249       }
250     }
251   }
252
253   void add(NodeType* node) {
254     std::lock_guard<MicroSpinLock> g(lock_);
255     if (nodes_.get() == nullptr) {
256       nodes_.reset(new std::vector<NodeType*>(1, node));
257     } else {
258       nodes_->push_back(node);
259     }
260     DCHECK_GT(refs(), 0);
261     dirty_.store(true, std::memory_order_relaxed);
262   }
263
264   int addRef() {
265     return refs_.fetch_add(1, std::memory_order_relaxed);
266   }
267
268   int releaseRef() {
269     // We don't expect to clean the recycler immediately everytime it is OK
270     // to do so. Here, it is possible that multiple accessors all release at
271     // the same time but nobody would clean the recycler here. If this
272     // happens, the recycler will usually still get cleaned when
273     // such a race doesn't happen. The worst case is the recycler will
274     // eventually get deleted along with the skiplist.
275     if (LIKELY(!dirty_.load(std::memory_order_relaxed) || refs() > 1)) {
276       return refs_.fetch_add(-1, std::memory_order_relaxed);
277     }
278
279     std::unique_ptr<std::vector<NodeType*>> newNodes;
280     {
281       std::lock_guard<MicroSpinLock> g(lock_);
282       if (nodes_.get() == nullptr || refs() > 1) {
283         return refs_.fetch_add(-1, std::memory_order_relaxed);
284       }
285       // once refs_ reaches 1 and there is no other accessor, it is safe to
286       // remove all the current nodes in the recycler, as we already acquired
287       // the lock here so no more new nodes can be added, even though new
288       // accessors may be added after that.
289       newNodes.swap(nodes_);
290       dirty_.store(false, std::memory_order_relaxed);
291     }
292
293     // TODO(xliu) should we spawn a thread to do this when there are large
294     // number of nodes in the recycler?
295     for (auto& node : *newNodes) {
296       NodeType::destroy(alloc_, node);
297     }
298
299     // decrease the ref count at the very end, to minimize the
300     // chance of other threads acquiring lock_ to clear the deleted
301     // nodes again.
302     return refs_.fetch_add(-1, std::memory_order_relaxed);
303   }
304
305   NodeAlloc& alloc() { return alloc_; }
306
307  private:
308   int refs() const {
309     return refs_.load(std::memory_order_relaxed);
310   }
311
312   std::unique_ptr<std::vector<NodeType*>> nodes_;
313   std::atomic<int32_t> refs_; // current number of visitors to the list
314   std::atomic<bool> dirty_; // whether *nodes_ is non-empty
315   MicroSpinLock lock_; // protects access to *nodes_
316   NodeAlloc alloc_;
317 };
318
319 // In case of arena allocator, no recycling is necessary, and it's possible
320 // to save on ConcurrentSkipList size.
321 template <typename NodeType, typename NodeAlloc>
322 class NodeRecycler<NodeType, NodeAlloc, typename std::enable_if<
323   NodeType::template DestroyIsNoOp<NodeAlloc>::value>::type> {
324  public:
325   explicit NodeRecycler(const NodeAlloc& alloc) : alloc_(alloc) { }
326
327   void addRef() { }
328   void releaseRef() { }
329
330   void add(NodeType* /* node */) {}
331
332   NodeAlloc& alloc() { return alloc_; }
333
334  private:
335   NodeAlloc alloc_;
336 };
337
338 }}  // namespaces