folly/experimental/RCURefCount.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 #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     std::array<RCURefCount*, 1> ptrs{{this}};
 101     useGlobal(ptrs);
 102   }
 103
 104   template <typename Container>
 105   static void useGlobal(const Container& refCountPtrs) {
 106     for (auto refCountPtr : refCountPtrs) {
 107       refCountPtr->state_ = State::GLOBAL_TRANSITION;
 108     }
 109
 110     synchronize_rcu();
 111     // At this point everyone is using the global count
 112
 113     for (auto refCountPtr : refCountPtrs) {
 114       auto accessor = refCountPtr->localCount_.accessAllThreads();
 115       for (auto& count : accessor) {
 116         count.collect();
 117       }
 118
 119       refCountPtr->state_ = State::GLOBAL;
 120     }
 121
 122     synchronize_rcu();
 123     // After this ++ or -- can return 0.
 124   }
 125
 126  private:
 127   using AtomicInt = std::atomic<Int>;
 128
 129   enum class State {
 130     LOCAL,
 131     GLOBAL_TRANSITION,
 132     GLOBAL
 133   };
 134
 135   class LocalRefCount {
 136    public:
 137     explicit LocalRefCount(AtomicInt& globalCount) :
 138         count_(0),
 139         globalCount_(globalCount) {
 140       RCURegisterThread();
 141     }
 142
 143     ~LocalRefCount() {
 144       collect();
 145     }
 146
 147     void collect() {
 148       globalCount_ += count_;
 149       count_ = 0;
 150     }
 151
 152     void operator++() {
 153       ++count_;
 154     }
 155
 156     void operator--() {
 157       --count_;
 158     }
 159
 160    private:
 161     Int count_;
 162     AtomicInt& globalCount_;
 163   };
 164
 165   std::atomic<State> state_{State::LOCAL};
 166   folly::ThreadLocal<LocalRefCount, RCURefCount> localCount_;
 167   std::atomic<int64_t> globalCount_{1};
 168 };
 169
 170 }