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 }