folly/concurrency/test/AtomicSharedPtrCounted.h

   1 /*
   2  * Copyright 2017-present 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 struct counted_shared_tag {};
  19 template <template <typename> class Atom = std::atomic>
  20 struct intrusive_shared_count {
  21   intrusive_shared_count() {
  22     counts.store(0);
  23   }
  24   void add_ref(uint64_t count = 1) {
  25     counts.fetch_add(count);
  26   }
  27
  28   uint64_t release_ref(uint64_t count = 1) {
  29     return counts.fetch_sub(count);
  30   }
  31   Atom<uint64_t> counts;
  32 };
  33
  34 template <template <typename> class Atom = std::atomic>
  35 struct counted_ptr_base {
  36  protected:
  37   static intrusive_shared_count<Atom>* getRef(void* pt) {
  38     char* p = (char*)pt;
  39     p -= sizeof(intrusive_shared_count<Atom>);
  40     return (intrusive_shared_count<Atom>*)p;
  41   }
  42 };
  43
  44 // basically shared_ptr, but only supports make_counted, and provides
  45 // access to add_ref / release_ref with a count.  Alias not supported.
  46 template <typename T, template <typename> class Atom = std::atomic>
  47 class counted_ptr : public counted_ptr_base<Atom> {
  48  public:
  49   T* p_;
  50   counted_ptr() : p_(nullptr) {}
  51   counted_ptr(counted_shared_tag, T* p) : p_(p) {
  52     if (p_)
  53       counted_ptr_base<Atom>::getRef(p_)->add_ref();
  54   }
  55
  56   counted_ptr(const counted_ptr& o) : p_(o.p_) {
  57     if (p_)
  58       counted_ptr_base<Atom>::getRef(p_)->add_ref();
  59   }
  60   counted_ptr& operator=(const counted_ptr& o) {
  61     if (p_ && counted_ptr_base<Atom>::getRef(p_)->release_ref() == 1) {
  62       p_->~T();
  63       free(counted_ptr_base<Atom>::getRef(p_));
  64     }
  65     p_ = o.p_;
  66     if (p_)
  67       counted_ptr_base<Atom>::getRef(p_)->add_ref();
  68     return *this;
  69   }
  70   explicit counted_ptr(T* p) : p_(p) {
  71     CHECK(!p);
  72   }
  73   ~counted_ptr() {
  74     if (p_ && counted_ptr_base<Atom>::getRef(p_)->release_ref() == 1) {
  75       p_->~T();
  76       free(counted_ptr_base<Atom>::getRef(p_));
  77     }
  78   }
  79   typename std::add_lvalue_reference<T>::type operator*() const {
  80     return *p_;
  81   }
  82
  83   T* get() const {
  84     return p_;
  85   }
  86   T* operator->() const {
  87     return p_;
  88   }
  89   explicit operator bool() const {
  90     return p_ == nullptr ? false : true;
  91   }
  92   bool operator==(const counted_ptr<T, Atom>& p) const {
  93     return get() == p.get();
  94   }
  95 };
  96
  97 template <
  98     template <typename> class Atom = std::atomic,
  99     typename T,
 100     typename... Args>
 101 counted_ptr<T, Atom> make_counted(Args&&... args) {
 102   char* mem = (char*)malloc(sizeof(T) + sizeof(intrusive_shared_count<Atom>));
 103   if (!mem) {
 104     throw std::bad_alloc();
 105   }
 106   new (mem) intrusive_shared_count<Atom>();
 107   T* ptr = (T*)(mem + sizeof(intrusive_shared_count<Atom>));
 108   new (ptr) T(std::forward<Args>(args)...);
 109   return counted_ptr<T, Atom>(counted_shared_tag(), ptr);
 110 }
 111
 112 template <template <typename> class Atom = std::atomic>
 113 class counted_ptr_internals : public counted_ptr_base<Atom> {
 114  public:
 115   template <typename T, typename... Args>
 116   static counted_ptr<T, Atom> make_ptr(Args&&... args) {
 117     return make_counted<Atom, T>(std::forward<Args...>(args...));
 118   }
 119   template <typename T>
 120   using CountedPtr = counted_ptr<T, Atom>;
 121   typedef void counted_base;
 122
 123   template <typename T>
 124   static counted_base* get_counted_base(const counted_ptr<T, Atom>& bar) {
 125     return bar.p_;
 126   }
 127
 128   template <typename T>
 129   static T* get_shared_ptr(counted_base* base) {
 130     return (T*)base;
 131   }
 132
 133   template <typename T>
 134   static T* release_ptr(counted_ptr<T, Atom>& p) {
 135     auto res = p.p_;
 136     p.p_ = nullptr;
 137     return res;
 138   }
 139
 140   template <typename T>
 141   static counted_ptr<T, Atom> get_shared_ptr_from_counted_base(
 142       counted_base* base,
 143       bool inc = true) {
 144     auto res = counted_ptr<T, Atom>(counted_shared_tag(), (T*)(base));
 145     if (!inc) {
 146       release_shared<T>(base, 1);
 147     }
 148     return res;
 149   }
 150
 151   static void inc_shared_count(counted_base* base, int64_t count) {
 152     counted_ptr_base<Atom>::getRef(base)->add_ref(count);
 153   }
 154
 155   template <typename T>
 156   static void release_shared(counted_base* base, uint64_t count) {
 157     if (count == counted_ptr_base<Atom>::getRef(base)->release_ref(count)) {
 158       ((T*)base)->~T();
 159       free(counted_ptr_base<Atom>::getRef(base));
 160     }
 161   }
 162 };