folly/experimental/RCURefCount.h

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